| Автор | SHA1 | Повідомлення | Дата |
|---|---|---|---|
|
|
4533c24fa1 | merge: dynamic video asset channel selection | 1 місяць тому |
|
|
f6321b3910 |
feat(channel): support dynamic video asset channel selection
Co-Authored-By: Codex <noreply@anthropic.com> |
1 місяць тому |
|
|
3fbf48313a |
feat(pricing): support usage billing and endpoint fixes
Add KlingAiping usage-token settlement from upstream final_unit_deduction and duration fallback handling. Normalize model endpoint config persistence and convert model pricing values between stored USD and display currency in the editor. Co-Authored-By: Codex <noreply@anthropic.com> |
1 місяць тому |
|
|
dd14839996 | merge: tianyiyun seedance channel | 1 місяць тому |
|
|
678add297f |
feat(channel): add tianyiyun seedance channel
Co-Authored-By: Codex <noreply@anthropic.com> |
1 місяць тому |
|
|
f100f9e409 | docs: design tianyiyun seedance channel | 1 місяць тому |
|
|
ffb3991654 | Merge branch 'feat/kling-aiping' | 1 місяць тому |
|
|
49ca8a6ae7 |
fix(kling): 修复原生路由和展示价格保存
- 避免仅编辑展示价格时写入默认真实计费规则 - 将 video-extend 纳入 KlingAiping 原生渠道约束 - 补充相关保存决策、渠道约束和计费展示测试 Co-Authored-By: Codex <noreply@anthropic.com> |
1 місяць тому |
|
|
6f805a6177 |
refactor: update aiping doubao upstream API paths to multimodal/sd endpoints
- Asset: /api/v1/volcengine/asset → /api/v1/multimodal/sd/assets
- Create video task: /api/v1/videos → /api/v1/multimodal/sd/videos/contents/generations/tasks
- Query video task: /api/v1/videos/{id} → /api/v1/multimodal/sd/videos/contents/generations/tasks/{id}
Co-Authored-By: Claude <noreply@anthropic.com>
|
1 місяць тому |
|
|
56400b6c61 |
feat: KlingAiping 渠道适配 + 展示价格管理后台
- 新增 KlingAiping 渠道 (ChannelTypeKlingAiping = 59) - 新增 KlingAiping 原生路由 (text2video/image2video/omni-video 等) - 新增 KlingAiping adaptor (relay/channel/task/kling/aiping/) - 前端模型定价页面新增「真实计费 / 展示价格」Tab 切换 - 新增 DisplayPricing 组件,支持模型展示价格的增删改 - 新增 displayPricingConfig 工具函数及测试 - 修复 KlingAiping 路由路径匹配与 task 错误响应 Co-Authored-By: Claude <noreply@anthropic.com> |
1 місяць тому |
|
|
a61ca078fc | fix: localize display pricing units | 1 місяць тому |
|
|
fee9504107 | feat: render display pricing on pricing page | 1 місяць тому |
|
|
cc97834e5b | feat: add display pricing frontend helpers | 1 місяць тому |
|
|
ced4d7f323 | feat: expose model display pricing | 1 місяць тому |
|
|
de117ac932 | feat: add model display pricing admin api | 1 місяць тому |
|
|
18a3495968 | feat: add model display pricing settings | 1 місяць тому |
|
|
302987989b | docs: expand display pricing frontend spec | 1 місяць тому |
|
|
6dd8c19ea8 | docs: clarify model display pricing design | 1 місяць тому |
|
|
6dd9bc6dec | docs: add model display pricing design | 1 місяць тому |
|
|
c9988af52b |
feat: 用户创建同步到从节点及 i18n 翻译
新建用户时异步同步到海外 slave 节点,确保跨区域用户数据一致。 补充前两个 commit 全部新功能的中英文翻译字符串。 Co-Authored-By: Claude <noreply@anthropic.com> |
1 місяць тому |
|
|
84d364d76a |
feat: 使用日志计费详情展开与 completion_tokens 修复
前端新增 UsageLogExpandedDetail 点击展开面板和 BillingDetailSummary hover 计费详情,后端修复 per_1m_tokens 任务结算日志的 completion_tokens 从零值改为写入 total_tokens 实际值。新增 taskBillingLogMerge 合并 任务预扣/结算日志。 Co-Authored-By: Claude <noreply@anthropic.com> |
1 місяць тому |
|
|
2e40bab478 |
feat: 模型定价矩阵价格取中位数展示
前端模型定价配置 UI 增强(Editor/GeneratorModal/PreviewPanel/PricingGrid), 后端新增 representativeMatrixPrice 取矩阵定价表中位数作为展示价格, 替代原来的 model_price=0 导致的价格展示为空问题。 Co-Authored-By: Claude <noreply@anthropic.com> |
1 місяць тому |
|
|
ec7c76f221 |
merge: feat/video-pricing-table-codex into master
- feat: video multidimensional matrix pricing system - feat: doubao aiping video channel support - feat: relay capture JSON storage with header redaction - feat: SSE stream error handling (Gemini/OpenAI/Responses) - test: E2E test infrastructure (Playwright + Chromium) Co-Authored-By: Claude <noreply@anthropic.com> |
1 місяць тому |
|
|
0e3b98efe0 |
chore: commit relay SSE error handling fixes and E2E mock server
- relay: SSE stream error exposure improvements (Claude/Gemini/OpenAI) - test: add mock capture server for E2E relay capture testing Co-Authored-By: Claude <noreply@anthropic.com> |
1 місяць тому |
|
|
139b6a6075 | feat: store relay capture records as json | 1 місяць тому |
|
|
65551e8003 | refactor: prepare relay capture json helpers | 1 місяць тому |
|
|
19d55dc255 | docs: design relay capture json storage | 1 місяць тому |
|
|
1d01c4cd64 |
feat: 增强流式响应错误处理,SSE内嵌错误正确曝光为API错误
- Gemini stream: 捕获 promptFeedback.blockReason 后停止流并返回 OpenAPI 错误 - OpenAI stream: 检测 SSE 帧中的 error 响应,在首帧写入前转换为 API 错误 - Responses stream: 处理 response.error / response.failed / error 类型帧 - 新增 ClearEventStreamHeadersIfNotWritten 工具方法,错误转换时清除流式头 - 新增 streamOpenAIErrorFromData,从 SSE data 帧提取 OpenAPI 错误结构 - 上游错误体会截断后存入 UpstreamBody 供排障 Co-Authored-By: Claude <noreply@anthropic.com> |
1 місяць тому |
|
|
696c549a66 |
feat(video): add doubao aiping video pricing support
Co-Authored-By: Codex <noreply@anthropic.com> |
1 місяць тому |
|
|
99168a9660 | test: stabilize model pricing e2e runtime | 1 місяць тому |
|
|
aeff58b122 | test: stabilize model pricing generator e2e | 1 місяць тому |
|
|
5451d801f3 | test: cover model pricing e2e workflow | 1 місяць тому |
|
|
a4b824952f | test: make model pricing tab selector clickable | 1 місяць тому |
|
|
32014dde92 | test: add model pricing e2e selectors | 1 місяць тому |
|
|
eaae02e767 |
test: ignore e2e artifacts
Co-Authored-By: Codex <noreply@anthropic.com> |
1 місяць тому |
|
|
75285bcabd |
test: align playwright runtime with chromium
Co-Authored-By: Codex <noreply@anthropic.com> |
1 місяць тому |
|
|
b976b486af |
test: avoid shell by default in e2e process helper
Co-Authored-By: Codex <noreply@anthropic.com> |
1 місяць тому |
|
|
568a91dfcf |
test: harden e2e helper error handling
Co-Authored-By: Codex <noreply@anthropic.com> |
1 місяць тому |
|
|
eb285055ac | test: harden e2e process cleanup | 1 місяць тому |
|
|
32603ba385 |
docs(test): video pricing table test plan with mock server design
Co-Authored-By: Claude <noreply@anthropic.com> |
1 місяць тому |
|
|
b91d07b0da | test: add e2e process helpers | 1 місяць тому |
|
|
9837e80fac | test: add e2e scripts | 1 місяць тому |
|
|
cf2c47b826 | docs: add e2e test foundation plan | 1 місяць тому |
|
|
915efefaec | docs: add e2e test foundation design | 1 місяць тому |
|
|
3ac94e2e6f | test(task-pricing): cover model pricing api and remix snapshots | 2 місяці тому |
|
|
19bdae5509 | fix(stream): default non-positive streaming timeout | 2 місяці тому |
|
|
089bc88015 | docs(task-pricing): document video pricing table testing | 2 місяці тому |
|
|
43b0bf293d | test(task-pricing): cover video usage billing flows | 2 місяці тому |
|
|
b36edf7f43 | feat(ui): add model multidimensional pricing editor | 2 місяці тому |
|
|
aff3284d02 | feat(ui): add model pricing config utilities | 2 місяці тому |
|
|
f3778e3931 | test(task-pricing): cover doubao usage parsing | 2 місяці тому |
|
|
285323d765 | feat(task-pricing): expose model pricing rules api | 2 місяці тому |
|
|
3a52f31608 | feat(task-pricing): gate channels for matrix usage billing | 2 місяці тому |
|
|
db8eacd558 | feat(task-pricing): persist matrix billing snapshots | 2 місяці тому |
|
|
8f2b014366 | feat(task-pricing): add matrix lookup decisions | 2 місяці тому |
|
|
ae1e68792a | feat(task-pricing): resolve pricing dimensions from task requests | 2 місяці тому |
|
|
b4301a53ae | feat(task-pricing): validate and cache model pricing rules | 2 місяці тому |
|
|
95060ea0aa | feat(task-pricing): add pricing config and decision types | 2 місяці тому |
| @@ -28,6 +28,19 @@ CLAUDE.md | |||
| logs/ | |||
| docs/superpowers | |||
| # Runtime request/response capture files | |||
| /[0-9]*.json | |||
| # Local dev server logs | |||
| web/vite-*.log | |||
| # E2E / Playwright artifacts | |||
| test-artifacts/ | |||
| web/.playwright/ | |||
| web/playwright-report/ | |||
| web/test-results/ | |||
| web/e2e/.auth/ | |||
| electron/node_modules | |||
| electron/dist | |||
| data/ | |||
| @@ -53,6 +53,10 @@ func ChannelType2APIType(channelType int) (int, bool) { | |||
| apiType = constant.APITypeMokaAI | |||
| case constant.ChannelTypeVolcEngine: | |||
| apiType = constant.APITypeVolcEngine | |||
| case constant.ChannelTypeDoubaoVideoCompatibleAiping: | |||
| apiType = constant.APITypeVolcEngine | |||
| case constant.ChannelTypeDoubaoVideoCompatibleTianyiYun: | |||
| apiType = constant.APITypeVolcEngine | |||
| case constant.ChannelTypeBaiduV2: | |||
| apiType = constant.APITypeBaiduV2 | |||
| case constant.ChannelTypeOpenRouter: | |||
| @@ -25,6 +25,7 @@ var defaultEndpointInfoMap = map[constant.EndpointType]EndpointInfo{ | |||
| constant.EndpointTypeJinaRerank: {Path: "/v1/rerank", Method: "POST"}, | |||
| constant.EndpointTypeImageGeneration: {Path: "/v1/images/generations", Method: "POST"}, | |||
| constant.EndpointTypeEmbeddings: {Path: "/v1/embeddings", Method: "POST"}, | |||
| constant.EndpointTypeDoubaoVideo: {Path: "/api/v3/contents/generations/tasks", Method: "POST"}, | |||
| } | |||
| // GetDefaultEndpointInfo 返回指定端点类型的默认信息以及是否存在 | |||
| @@ -30,6 +30,8 @@ func GetEndpointTypesByChannelType(channelType int, modelName string) []constant | |||
| endpointTypes = []constant.EndpointType{constant.EndpointTypeOpenAI, constant.EndpointTypeOpenAIResponse} | |||
| case constant.ChannelTypeSora: | |||
| endpointTypes = []constant.EndpointType{constant.EndpointTypeOpenAIVideo} | |||
| case constant.ChannelTypeDoubaoVideo, constant.ChannelTypeDoubaoVideoCompatibleAiping, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun: | |||
| endpointTypes = []constant.EndpointType{constant.EndpointTypeDoubaoVideo} | |||
| default: | |||
| if IsOpenAIResponseOnlyModel(modelName) { | |||
| endpointTypes = []constant.EndpointType{constant.EndpointTypeOpenAIResponse} | |||
| @@ -0,0 +1,21 @@ | |||
| package common | |||
| import ( | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/constant" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestTianyiYunChannelUsesVolcEngineAPIType(t *testing.T) { | |||
| apiType, ok := ChannelType2APIType(constant.ChannelTypeDoubaoVideoCompatibleTianyiYun) | |||
| require.True(t, ok) | |||
| require.Equal(t, constant.APITypeVolcEngine, apiType) | |||
| } | |||
| func TestTianyiYunChannelUsesDoubaoVideoEndpointType(t *testing.T) { | |||
| got := GetEndpointTypesByChannelType(constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "cdance2.0-0611") | |||
| require.Equal(t, []constant.EndpointType{constant.EndpointTypeDoubaoVideo}, got) | |||
| } | |||
| @@ -276,6 +276,10 @@ func Max(a int, b int) int { | |||
| } | |||
| func MessageWithRequestId(message string, id string) string { | |||
| id = strings.TrimSpace(id) | |||
| if id == "" || strings.Contains(message, "request id:") { | |||
| return message | |||
| } | |||
| return fmt.Sprintf("%s (request id: %s)", message, id) | |||
| } | |||
| @@ -0,0 +1,27 @@ | |||
| package common | |||
| import "testing" | |||
| func TestMessageWithRequestIdAppendsRequestId(t *testing.T) { | |||
| got := MessageWithRequestId("upstream failed", "req-123") | |||
| want := "upstream failed (request id: req-123)" | |||
| if got != want { | |||
| t.Fatalf("MessageWithRequestId() = %q, want %q", got, want) | |||
| } | |||
| } | |||
| func TestMessageWithRequestIdSkipsEmptyRequestId(t *testing.T) { | |||
| got := MessageWithRequestId("upstream failed", " ") | |||
| want := "upstream failed" | |||
| if got != want { | |||
| t.Fatalf("MessageWithRequestId() = %q, want %q", got, want) | |||
| } | |||
| } | |||
| func TestMessageWithRequestIdDoesNotAppendTwice(t *testing.T) { | |||
| message := "upstream failed (request id: req-123)" | |||
| got := MessageWithRequestId(message, "req-123") | |||
| if got != message { | |||
| t.Fatalf("MessageWithRequestId() = %q, want %q", got, message) | |||
| } | |||
| } | |||
| @@ -1,61 +1,64 @@ | |||
| package constant | |||
| const ( | |||
| ChannelTypeUnknown = 0 | |||
| ChannelTypeOpenAI = 1 | |||
| ChannelTypeMidjourney = 2 | |||
| ChannelTypeAzure = 3 | |||
| ChannelTypeOllama = 4 | |||
| ChannelTypeMidjourneyPlus = 5 | |||
| ChannelTypeOpenAIMax = 6 | |||
| ChannelTypeOhMyGPT = 7 | |||
| ChannelTypeCustom = 8 | |||
| ChannelTypeAILS = 9 | |||
| ChannelTypeAIProxy = 10 | |||
| ChannelTypePaLM = 11 | |||
| ChannelTypeAPI2GPT = 12 | |||
| ChannelTypeAIGC2D = 13 | |||
| ChannelTypeAnthropic = 14 | |||
| ChannelTypeBaidu = 15 | |||
| ChannelTypeZhipu = 16 | |||
| ChannelTypeAli = 17 | |||
| ChannelTypeXunfei = 18 | |||
| ChannelType360 = 19 | |||
| ChannelTypeOpenRouter = 20 | |||
| ChannelTypeAIProxyLibrary = 21 | |||
| ChannelTypeFastGPT = 22 | |||
| ChannelTypeTencent = 23 | |||
| ChannelTypeGemini = 24 | |||
| ChannelTypeMoonshot = 25 | |||
| ChannelTypeZhipu_v4 = 26 | |||
| ChannelTypePerplexity = 27 | |||
| ChannelTypeLingYiWanWu = 31 | |||
| ChannelTypeAws = 33 | |||
| ChannelTypeCohere = 34 | |||
| ChannelTypeMiniMax = 35 | |||
| ChannelTypeSunoAPI = 36 | |||
| ChannelTypeDify = 37 | |||
| ChannelTypeJina = 38 | |||
| ChannelCloudflare = 39 | |||
| ChannelTypeSiliconFlow = 40 | |||
| ChannelTypeVertexAi = 41 | |||
| ChannelTypeMistral = 42 | |||
| ChannelTypeDeepSeek = 43 | |||
| ChannelTypeMokaAI = 44 | |||
| ChannelTypeVolcEngine = 45 | |||
| ChannelTypeBaiduV2 = 46 | |||
| ChannelTypeXinference = 47 | |||
| ChannelTypeXai = 48 | |||
| ChannelTypeCoze = 49 | |||
| ChannelTypeKling = 50 | |||
| ChannelTypeJimeng = 51 | |||
| ChannelTypeVidu = 52 | |||
| ChannelTypeSubmodel = 53 | |||
| ChannelTypeDoubaoVideo = 54 | |||
| ChannelTypeSora = 55 | |||
| ChannelTypeReplicate = 56 | |||
| ChannelTypeCodex = 57 | |||
| ChannelTypeDummy // this one is only for count, do not add any channel after this | |||
| ChannelTypeUnknown = 0 | |||
| ChannelTypeOpenAI = 1 | |||
| ChannelTypeMidjourney = 2 | |||
| ChannelTypeAzure = 3 | |||
| ChannelTypeOllama = 4 | |||
| ChannelTypeMidjourneyPlus = 5 | |||
| ChannelTypeOpenAIMax = 6 | |||
| ChannelTypeOhMyGPT = 7 | |||
| ChannelTypeCustom = 8 | |||
| ChannelTypeAILS = 9 | |||
| ChannelTypeAIProxy = 10 | |||
| ChannelTypePaLM = 11 | |||
| ChannelTypeAPI2GPT = 12 | |||
| ChannelTypeAIGC2D = 13 | |||
| ChannelTypeAnthropic = 14 | |||
| ChannelTypeBaidu = 15 | |||
| ChannelTypeZhipu = 16 | |||
| ChannelTypeAli = 17 | |||
| ChannelTypeXunfei = 18 | |||
| ChannelType360 = 19 | |||
| ChannelTypeOpenRouter = 20 | |||
| ChannelTypeAIProxyLibrary = 21 | |||
| ChannelTypeFastGPT = 22 | |||
| ChannelTypeTencent = 23 | |||
| ChannelTypeGemini = 24 | |||
| ChannelTypeMoonshot = 25 | |||
| ChannelTypeZhipu_v4 = 26 | |||
| ChannelTypePerplexity = 27 | |||
| ChannelTypeLingYiWanWu = 31 | |||
| ChannelTypeAws = 33 | |||
| ChannelTypeCohere = 34 | |||
| ChannelTypeMiniMax = 35 | |||
| ChannelTypeSunoAPI = 36 | |||
| ChannelTypeDify = 37 | |||
| ChannelTypeJina = 38 | |||
| ChannelCloudflare = 39 | |||
| ChannelTypeSiliconFlow = 40 | |||
| ChannelTypeVertexAi = 41 | |||
| ChannelTypeMistral = 42 | |||
| ChannelTypeDeepSeek = 43 | |||
| ChannelTypeMokaAI = 44 | |||
| ChannelTypeVolcEngine = 45 | |||
| ChannelTypeBaiduV2 = 46 | |||
| ChannelTypeXinference = 47 | |||
| ChannelTypeXai = 48 | |||
| ChannelTypeCoze = 49 | |||
| ChannelTypeKling = 50 | |||
| ChannelTypeJimeng = 51 | |||
| ChannelTypeVidu = 52 | |||
| ChannelTypeSubmodel = 53 | |||
| ChannelTypeDoubaoVideo = 54 | |||
| ChannelTypeSora = 55 | |||
| ChannelTypeReplicate = 56 | |||
| ChannelTypeCodex = 57 | |||
| ChannelTypeDoubaoVideoCompatibleAiping = 58 | |||
| ChannelTypeKlingAiping = 59 | |||
| ChannelTypeDoubaoVideoCompatibleTianyiYun = 60 | |||
| ChannelTypeDummy // this one is only for count, do not add any channel after this | |||
| ) | |||
| @@ -118,63 +121,69 @@ var ChannelBaseURLs = []string{ | |||
| "https://api.openai.com", //55 | |||
| "https://api.replicate.com", //56 | |||
| "https://chatgpt.com", //57 | |||
| "", //58 | |||
| "https://aiping.cn/api", //59 | |||
| "https://ai.ctaigw.cn", //60 | |||
| } | |||
| var ChannelTypeNames = map[int]string{ | |||
| ChannelTypeUnknown: "Unknown", | |||
| ChannelTypeOpenAI: "OpenAI", | |||
| ChannelTypeMidjourney: "Midjourney", | |||
| ChannelTypeAzure: "Azure", | |||
| ChannelTypeOllama: "Ollama", | |||
| ChannelTypeMidjourneyPlus: "MidjourneyPlus", | |||
| ChannelTypeOpenAIMax: "OpenAIMax", | |||
| ChannelTypeOhMyGPT: "OhMyGPT", | |||
| ChannelTypeCustom: "Custom", | |||
| ChannelTypeAILS: "AILS", | |||
| ChannelTypeAIProxy: "AIProxy", | |||
| ChannelTypePaLM: "PaLM", | |||
| ChannelTypeAPI2GPT: "API2GPT", | |||
| ChannelTypeAIGC2D: "AIGC2D", | |||
| ChannelTypeAnthropic: "Anthropic", | |||
| ChannelTypeBaidu: "Baidu", | |||
| ChannelTypeZhipu: "Zhipu", | |||
| ChannelTypeAli: "Ali", | |||
| ChannelTypeXunfei: "Xunfei", | |||
| ChannelType360: "360", | |||
| ChannelTypeOpenRouter: "OpenRouter", | |||
| ChannelTypeAIProxyLibrary: "AIProxyLibrary", | |||
| ChannelTypeFastGPT: "FastGPT", | |||
| ChannelTypeTencent: "Tencent", | |||
| ChannelTypeGemini: "Gemini", | |||
| ChannelTypeMoonshot: "Moonshot", | |||
| ChannelTypeZhipu_v4: "ZhipuV4", | |||
| ChannelTypePerplexity: "Perplexity", | |||
| ChannelTypeLingYiWanWu: "LingYiWanWu", | |||
| ChannelTypeAws: "AWS", | |||
| ChannelTypeCohere: "Cohere", | |||
| ChannelTypeMiniMax: "MiniMax", | |||
| ChannelTypeSunoAPI: "SunoAPI", | |||
| ChannelTypeDify: "Dify", | |||
| ChannelTypeJina: "Jina", | |||
| ChannelCloudflare: "Cloudflare", | |||
| ChannelTypeSiliconFlow: "SiliconFlow", | |||
| ChannelTypeVertexAi: "VertexAI", | |||
| ChannelTypeMistral: "Mistral", | |||
| ChannelTypeDeepSeek: "DeepSeek", | |||
| ChannelTypeMokaAI: "MokaAI", | |||
| ChannelTypeVolcEngine: "VolcEngine", | |||
| ChannelTypeBaiduV2: "BaiduV2", | |||
| ChannelTypeXinference: "Xinference", | |||
| ChannelTypeXai: "xAI", | |||
| ChannelTypeCoze: "Coze", | |||
| ChannelTypeKling: "Kling", | |||
| ChannelTypeJimeng: "Jimeng", | |||
| ChannelTypeVidu: "Vidu", | |||
| ChannelTypeSubmodel: "Submodel", | |||
| ChannelTypeDoubaoVideo: "DoubaoVideo", | |||
| ChannelTypeSora: "Sora", | |||
| ChannelTypeReplicate: "Replicate", | |||
| ChannelTypeCodex: "Codex", | |||
| ChannelTypeUnknown: "Unknown", | |||
| ChannelTypeOpenAI: "OpenAI", | |||
| ChannelTypeMidjourney: "Midjourney", | |||
| ChannelTypeAzure: "Azure", | |||
| ChannelTypeOllama: "Ollama", | |||
| ChannelTypeMidjourneyPlus: "MidjourneyPlus", | |||
| ChannelTypeOpenAIMax: "OpenAIMax", | |||
| ChannelTypeOhMyGPT: "OhMyGPT", | |||
| ChannelTypeCustom: "Custom", | |||
| ChannelTypeAILS: "AILS", | |||
| ChannelTypeAIProxy: "AIProxy", | |||
| ChannelTypePaLM: "PaLM", | |||
| ChannelTypeAPI2GPT: "API2GPT", | |||
| ChannelTypeAIGC2D: "AIGC2D", | |||
| ChannelTypeAnthropic: "Anthropic", | |||
| ChannelTypeBaidu: "Baidu", | |||
| ChannelTypeZhipu: "Zhipu", | |||
| ChannelTypeAli: "Ali", | |||
| ChannelTypeXunfei: "Xunfei", | |||
| ChannelType360: "360", | |||
| ChannelTypeOpenRouter: "OpenRouter", | |||
| ChannelTypeAIProxyLibrary: "AIProxyLibrary", | |||
| ChannelTypeFastGPT: "FastGPT", | |||
| ChannelTypeTencent: "Tencent", | |||
| ChannelTypeGemini: "Gemini", | |||
| ChannelTypeMoonshot: "Moonshot", | |||
| ChannelTypeZhipu_v4: "ZhipuV4", | |||
| ChannelTypePerplexity: "Perplexity", | |||
| ChannelTypeLingYiWanWu: "LingYiWanWu", | |||
| ChannelTypeAws: "AWS", | |||
| ChannelTypeCohere: "Cohere", | |||
| ChannelTypeMiniMax: "MiniMax", | |||
| ChannelTypeSunoAPI: "SunoAPI", | |||
| ChannelTypeDify: "Dify", | |||
| ChannelTypeJina: "Jina", | |||
| ChannelCloudflare: "Cloudflare", | |||
| ChannelTypeSiliconFlow: "SiliconFlow", | |||
| ChannelTypeVertexAi: "VertexAI", | |||
| ChannelTypeMistral: "Mistral", | |||
| ChannelTypeDeepSeek: "DeepSeek", | |||
| ChannelTypeMokaAI: "MokaAI", | |||
| ChannelTypeVolcEngine: "VolcEngine", | |||
| ChannelTypeBaiduV2: "BaiduV2", | |||
| ChannelTypeXinference: "Xinference", | |||
| ChannelTypeXai: "xAI", | |||
| ChannelTypeCoze: "Coze", | |||
| ChannelTypeKling: "Kling", | |||
| ChannelTypeJimeng: "Jimeng", | |||
| ChannelTypeVidu: "Vidu", | |||
| ChannelTypeSubmodel: "Submodel", | |||
| ChannelTypeDoubaoVideo: "DoubaoVideo", | |||
| ChannelTypeSora: "Sora", | |||
| ChannelTypeReplicate: "Replicate", | |||
| ChannelTypeCodex: "Codex", | |||
| ChannelTypeDoubaoVideoCompatibleAiping: "DoubaoVideoCompatibleAiping", | |||
| ChannelTypeKlingAiping: "KlingAiping", | |||
| ChannelTypeDoubaoVideoCompatibleTianyiYun: "DoubaoVideoCompatibleTianyiYun", | |||
| } | |||
| func GetChannelTypeName(channelType int) string { | |||
| @@ -24,6 +24,7 @@ const ( | |||
| EndpointTypeEmbeddings EndpointType = "embeddings" | |||
| // OpenAI Video API,如 Sora 视频生成 | |||
| EndpointTypeOpenAIVideo EndpointType = "openai-video" | |||
| EndpointTypeDoubaoVideo EndpointType = "doubao-video" | |||
| //EndpointTypeMidjourney EndpointType = "midjourney-proxy" | |||
| //EndpointTypeSuno EndpointType = "suno-proxy" | |||
| //EndpointTypeKling EndpointType = "kling" | |||
| @@ -65,6 +65,8 @@ func testChannel(channel *model.Channel, testModel string, endpointType string, | |||
| constant.ChannelTypeKling, | |||
| constant.ChannelTypeJimeng, | |||
| constant.ChannelTypeDoubaoVideo, | |||
| constant.ChannelTypeDoubaoVideoCompatibleAiping, | |||
| constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, | |||
| constant.ChannelTypeVidu, | |||
| } | |||
| if lo.Contains(unsupportedTestChannelTypes, channel.Type) { | |||
| @@ -0,0 +1,36 @@ | |||
| package controller | |||
| import ( | |||
| "strings" | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/constant" | |||
| "github.com/QuantumNous/new-api/model" | |||
| "github.com/QuantumNous/new-api/service" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestTestChannelRejectsAsyncVideoChannels(t *testing.T) { | |||
| channels := []*model.Channel{ | |||
| {Type: constant.ChannelTypeDoubaoVideoCompatibleAiping, Models: "doubao-seedance-2-0-260128", Status: common.ChannelStatusEnabled}, | |||
| {Type: constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, Models: "cdance2.0-0611", Status: common.ChannelStatusEnabled}, | |||
| } | |||
| for _, channel := range channels { | |||
| result := testChannel(channel, "", "", false) | |||
| require.Error(t, result.localErr) | |||
| require.True(t, strings.Contains(result.localErr.Error(), "channel test is not supported")) | |||
| } | |||
| } | |||
| func TestRequiredTaskChannelTypeForTianyiYunSeedanceModelUsesAllowedFamily(t *testing.T) { | |||
| c := newControllerJSONContext(t, "/api/v3/contents/generations/tasks", `{"model":"Doubao-Seedance-2.0"}`) | |||
| require.Equal(t, 0, requiredTaskChannelTypeForRequest(c)) | |||
| require.ElementsMatch(t, | |||
| service.VideoAssetChannelTypesForFamily(service.VideoAssetFamilySeedance), | |||
| allowedTaskChannelTypesForRequest(c), | |||
| ) | |||
| } | |||
| @@ -0,0 +1,158 @@ | |||
| package controller | |||
| import ( | |||
| "fmt" | |||
| "net/http" | |||
| "strings" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/constant" | |||
| "github.com/QuantumNous/new-api/dto" | |||
| "github.com/QuantumNous/new-api/model" | |||
| relaycommon "github.com/QuantumNous/new-api/relay/common" | |||
| "github.com/QuantumNous/new-api/types" | |||
| "github.com/gin-gonic/gin" | |||
| ) | |||
| func convertAipingNativeVideoRequest(body []byte) (relaycommon.TaskSubmitReq, error) { | |||
| var raw map[string]any | |||
| if err := common.Unmarshal(body, &raw); err != nil { | |||
| return relaycommon.TaskSubmitReq{}, err | |||
| } | |||
| modelName, _ := raw["model"].(string) | |||
| if strings.TrimSpace(modelName) == "" { | |||
| return relaycommon.TaskSubmitReq{}, fmt.Errorf("model field is required") | |||
| } | |||
| metadata := map[string]any{} | |||
| for key, value := range raw { | |||
| if key == "model" || key == "content" { | |||
| continue | |||
| } | |||
| metadata[key] = value | |||
| } | |||
| var passthroughContent []any | |||
| if rawContent, ok := raw["content"].([]any); ok { | |||
| for _, item := range rawContent { | |||
| itemMap, ok := item.(map[string]any) | |||
| if !ok { | |||
| passthroughContent = append(passthroughContent, item) | |||
| continue | |||
| } | |||
| passthroughContent = append(passthroughContent, itemMap) | |||
| } | |||
| } | |||
| if len(passthroughContent) > 0 { | |||
| metadata["content"] = passthroughContent | |||
| } | |||
| return relaycommon.TaskSubmitReq{ | |||
| Model: strings.TrimSpace(modelName), | |||
| Metadata: metadata, | |||
| }, nil | |||
| } | |||
| func prepareAipingNativeVideoSubmit(c *gin.Context, info *relaycommon.RelayInfo) error { | |||
| bodyStorage, err := common.GetBodyStorage(c) | |||
| if err != nil { | |||
| return err | |||
| } | |||
| body, err := bodyStorage.Bytes() | |||
| if err != nil { | |||
| return err | |||
| } | |||
| req, err := convertAipingNativeVideoRequest(body) | |||
| if err != nil { | |||
| return err | |||
| } | |||
| info.OriginModelName = req.Model | |||
| info.Action = constant.TaskActionGenerate | |||
| relaycommon.StoreTaskRequest(c, info, constant.TaskActionGenerate, req) | |||
| return nil | |||
| } | |||
| func AipingNativeVideoSubmit(c *gin.Context) { | |||
| relayInfo, err := relaycommon.GenRelayInfo(c, types.RelayFormatTask, nil, nil) | |||
| if err != nil { | |||
| c.JSON(http.StatusInternalServerError, &dto.TaskError{ | |||
| Code: "gen_relay_info_failed", | |||
| Message: err.Error(), | |||
| StatusCode: http.StatusInternalServerError, | |||
| }) | |||
| return | |||
| } | |||
| if relayInfo.ChannelMeta == nil { | |||
| relayInfo.ChannelMeta = &relaycommon.ChannelMeta{} | |||
| } | |||
| if err := prepareAipingNativeVideoSubmit(c, relayInfo); err != nil { | |||
| c.JSON(http.StatusBadRequest, &dto.TaskError{ | |||
| Code: "invalid_request", | |||
| Message: err.Error(), | |||
| StatusCode: http.StatusBadRequest, | |||
| }) | |||
| return | |||
| } | |||
| relayTaskWithInfo(c, relayInfo) | |||
| } | |||
| func buildAipingNativeFetchResponse(task *model.Task) ([]byte, error) { | |||
| payload := map[string]any{} | |||
| if len(task.Data) > 0 { | |||
| if err := common.Unmarshal(task.Data, &payload); err != nil { | |||
| return nil, err | |||
| } | |||
| } | |||
| delete(payload, "aiping_id") | |||
| if _, ok := payload["error"]; !ok { | |||
| payload["error"] = nil | |||
| } | |||
| payload["id"] = task.TaskID | |||
| if _, ok := payload["model"]; !ok { | |||
| payload["model"] = task.Properties.OriginModelName | |||
| } | |||
| if _, ok := payload["status"]; !ok { | |||
| payload["status"] = mapAipingNativeStatus(task.Status) | |||
| } | |||
| if _, ok := payload["created_at"]; !ok { | |||
| payload["created_at"] = task.CreatedAt | |||
| } | |||
| if _, ok := payload["updated_at"]; !ok { | |||
| payload["updated_at"] = task.UpdatedAt | |||
| } | |||
| return common.Marshal(payload) | |||
| } | |||
| func mapAipingNativeStatus(status model.TaskStatus) string { | |||
| switch status { | |||
| case model.TaskStatusQueued, model.TaskStatusSubmitted: | |||
| return "queued" | |||
| case model.TaskStatusInProgress: | |||
| return "running" | |||
| case model.TaskStatusSuccess: | |||
| return "succeeded" | |||
| case model.TaskStatusFailure: | |||
| return "failed" | |||
| default: | |||
| return "running" | |||
| } | |||
| } | |||
| func AipingNativeVideoFetch(c *gin.Context) { | |||
| taskID := c.Param("task_id") | |||
| task, exist, err := model.GetByTaskId(c.GetInt("id"), taskID) | |||
| if err != nil { | |||
| c.JSON(http.StatusInternalServerError, &dto.TaskError{Code: "get_task_failed", Message: err.Error(), StatusCode: http.StatusInternalServerError}) | |||
| return | |||
| } | |||
| if !exist { | |||
| c.JSON(http.StatusNotFound, gin.H{"error": gin.H{"code": "not_found", "message": "task not found"}}) | |||
| return | |||
| } | |||
| data, err := buildAipingNativeFetchResponse(task) | |||
| if err != nil { | |||
| c.JSON(http.StatusInternalServerError, &dto.TaskError{Code: "build_response_failed", Message: err.Error(), StatusCode: http.StatusInternalServerError}) | |||
| return | |||
| } | |||
| c.Data(http.StatusOK, "application/json", data) | |||
| } | |||
| @@ -0,0 +1,209 @@ | |||
| package controller | |||
| import ( | |||
| "net/http" | |||
| "net/http/httptest" | |||
| "strings" | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/constant" | |||
| "github.com/QuantumNous/new-api/model" | |||
| relaycommon "github.com/QuantumNous/new-api/relay/common" | |||
| "github.com/gin-gonic/gin" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestConvertAipingNativeVideoRequestPreservesNativeContentOrder(t *testing.T) { | |||
| req, err := convertAipingNativeVideoRequest([]byte(`{ | |||
| "model":"doubao-seedance-2-0-260128", | |||
| "content":[ | |||
| {"type":"text","text":"first prompt"}, | |||
| {"type":"image_url","image_url":{"url":"asset://img1"},"role":"reference_image"}, | |||
| {"type":"text","text":"second prompt"}, | |||
| {"type":"video_url","video_url":{"url":"https://example.test/input.mp4"},"role":"reference_video"} | |||
| ], | |||
| "duration":5, | |||
| "resolution":"480p", | |||
| "tools":[{"type":"web_search"}] | |||
| }`)) | |||
| require.NoError(t, err) | |||
| require.Equal(t, "doubao-seedance-2-0-260128", req.Model) | |||
| require.Empty(t, req.Prompt) | |||
| require.Equal(t, float64(5), req.Metadata["duration"]) | |||
| require.Equal(t, "480p", req.Metadata["resolution"]) | |||
| require.NotNil(t, req.Metadata["tools"]) | |||
| content, ok := req.Metadata["content"].([]any) | |||
| require.True(t, ok) | |||
| require.Len(t, content, 4) | |||
| require.Equal(t, "text", content[0].(map[string]any)["type"]) | |||
| require.Equal(t, "image_url", content[1].(map[string]any)["type"]) | |||
| require.Equal(t, "text", content[2].(map[string]any)["type"]) | |||
| require.Equal(t, "video_url", content[3].(map[string]any)["type"]) | |||
| } | |||
| func TestConvertAipingNativeVideoRequestPreservesArkDocumentFields(t *testing.T) { | |||
| req, err := convertAipingNativeVideoRequest([]byte(`{ | |||
| "model":"doubao-seedance-2-0-260128", | |||
| "content":[ | |||
| {"type":"text","text":"开场:海边日落"}, | |||
| {"type":"image_url","image_url":{"url":"asset://image-1"},"role":"reference_image"}, | |||
| {"type":"video_url","video_url":{"url":"asset://video-1"},"role":"reference_video"}, | |||
| {"type":"audio_url","audio_url":{"url":"asset://audio-1"},"role":"reference_audio"}, | |||
| {"type":"draft_task","draft_task":{"id":"cgt-draft"}} | |||
| ], | |||
| "callback_url":"https://example.test/callback", | |||
| "return_last_frame":true, | |||
| "service_tier":"default", | |||
| "execution_expires_after":3600, | |||
| "generate_audio":false, | |||
| "draft":true, | |||
| "tools":[{"type":"web_search"}], | |||
| "safety_identifier":"user-hash-1", | |||
| "priority":5, | |||
| "resolution":"480p", | |||
| "ratio":"1:1", | |||
| "duration":5, | |||
| "frames":29, | |||
| "seed":11, | |||
| "camera_fixed":false, | |||
| "watermark":true | |||
| }`)) | |||
| require.NoError(t, err) | |||
| require.Equal(t, "doubao-seedance-2-0-260128", req.Model) | |||
| require.Empty(t, req.Prompt) | |||
| require.Equal(t, "https://example.test/callback", req.Metadata["callback_url"]) | |||
| require.Equal(t, true, req.Metadata["return_last_frame"]) | |||
| require.Equal(t, "default", req.Metadata["service_tier"]) | |||
| require.Equal(t, float64(3600), req.Metadata["execution_expires_after"]) | |||
| require.Equal(t, false, req.Metadata["generate_audio"]) | |||
| require.Equal(t, true, req.Metadata["draft"]) | |||
| require.Equal(t, "user-hash-1", req.Metadata["safety_identifier"]) | |||
| require.Equal(t, float64(5), req.Metadata["priority"]) | |||
| require.Equal(t, "480p", req.Metadata["resolution"]) | |||
| require.Equal(t, "1:1", req.Metadata["ratio"]) | |||
| require.Equal(t, float64(5), req.Metadata["duration"]) | |||
| require.Equal(t, float64(29), req.Metadata["frames"]) | |||
| require.Equal(t, float64(11), req.Metadata["seed"]) | |||
| require.Equal(t, false, req.Metadata["camera_fixed"]) | |||
| require.Equal(t, true, req.Metadata["watermark"]) | |||
| content, ok := req.Metadata["content"].([]any) | |||
| require.True(t, ok) | |||
| require.Len(t, content, 5) | |||
| require.Equal(t, "audio_url", content[3].(map[string]any)["type"]) | |||
| require.Equal(t, "draft_task", content[4].(map[string]any)["type"]) | |||
| tools, ok := req.Metadata["tools"].([]any) | |||
| require.True(t, ok) | |||
| require.Equal(t, "web_search", tools[0].(map[string]any)["type"]) | |||
| } | |||
| func TestConvertAipingNativeVideoRequestRejectsMissingModel(t *testing.T) { | |||
| _, err := convertAipingNativeVideoRequest([]byte(`{"content":[{"type":"text","text":"prompt"}]}`)) | |||
| require.ErrorContains(t, err, "model") | |||
| } | |||
| func TestPrepareAipingNativeVideoSubmitStoresConvertedRequest(t *testing.T) { | |||
| c := newControllerJSONContext(t, "/api/v3/contents/generations/tasks", `{ | |||
| "model":"doubao-seedance-2-0-260128", | |||
| "content":[{"type":"text","text":"prompt"}] | |||
| }`) | |||
| info := &relaycommon.RelayInfo{TaskRelayInfo: &relaycommon.TaskRelayInfo{}} | |||
| err := prepareAipingNativeVideoSubmit(c, info) | |||
| require.NoError(t, err) | |||
| stored, err := relaycommon.GetTaskRequest(c) | |||
| require.NoError(t, err) | |||
| require.Equal(t, "doubao-seedance-2-0-260128", stored.Model) | |||
| require.Empty(t, stored.Prompt) | |||
| content, ok := stored.Metadata["content"].([]any) | |||
| require.True(t, ok) | |||
| require.Equal(t, "prompt", content[0].(map[string]any)["text"]) | |||
| require.Equal(t, constant.TaskActionGenerate, info.Action) | |||
| require.Equal(t, "doubao-seedance-2-0-260128", info.OriginModelName) | |||
| } | |||
| func TestBuildAipingNativeFetchResponseRewritesIDAndPreservesUpstreamFields(t *testing.T) { | |||
| task := &model.Task{ | |||
| TaskID: "task_public", | |||
| Status: model.TaskStatusSuccess, | |||
| CreatedAt: 1781496040, | |||
| UpdatedAt: 1781496278, | |||
| Properties: model.Properties{OriginModelName: "doubao-seedance-2-0-260128"}, | |||
| Data: []byte(`{ | |||
| "id":"cgt-upstream", | |||
| "aiping_id":"3d2c8c17-36ad-4138-8602-88b3c60e56c6", | |||
| "model":"doubao-seedance-2-0-260128", | |||
| "status":"succeeded", | |||
| "content":{"video_url":"https://example.test/output.mp4"}, | |||
| "usage":{"completion_tokens":48400,"total_tokens":48400}, | |||
| "created_at":1781496040, | |||
| "updated_at":1781496278, | |||
| "seed":73812, | |||
| "resolution":"480p", | |||
| "ratio":"1:1", | |||
| "duration":5, | |||
| "framespersecond":24, | |||
| "service_tier":"default", | |||
| "execution_expires_after":172800, | |||
| "generate_audio":true, | |||
| "draft":false, | |||
| "priority":0 | |||
| }`), | |||
| } | |||
| data, err := buildAipingNativeFetchResponse(task) | |||
| require.NoError(t, err) | |||
| var payload map[string]any | |||
| require.NoError(t, common.Unmarshal(data, &payload)) | |||
| require.Equal(t, "task_public", payload["id"]) | |||
| require.Equal(t, "succeeded", payload["status"]) | |||
| require.Equal(t, "doubao-seedance-2-0-260128", payload["model"]) | |||
| require.Equal(t, "https://example.test/output.mp4", payload["content"].(map[string]any)["video_url"]) | |||
| require.Equal(t, float64(48400), payload["usage"].(map[string]any)["total_tokens"]) | |||
| require.Equal(t, float64(73812), payload["seed"]) | |||
| require.Equal(t, "480p", payload["resolution"]) | |||
| require.Equal(t, "1:1", payload["ratio"]) | |||
| require.Equal(t, float64(5), payload["duration"]) | |||
| require.Equal(t, float64(24), payload["framespersecond"]) | |||
| require.Equal(t, "default", payload["service_tier"]) | |||
| require.Equal(t, float64(172800), payload["execution_expires_after"]) | |||
| require.Equal(t, true, payload["generate_audio"]) | |||
| require.Equal(t, false, payload["draft"]) | |||
| require.Equal(t, float64(0), payload["priority"]) | |||
| require.Contains(t, payload, "error") | |||
| require.Nil(t, payload["error"]) | |||
| require.NotContains(t, payload, "aiping_id") | |||
| } | |||
| func TestBuildAipingNativeFetchResponseFallsBackForFailedTask(t *testing.T) { | |||
| task := &model.Task{ | |||
| TaskID: "task_public", | |||
| Status: model.TaskStatusFailure, | |||
| CreatedAt: 100, | |||
| UpdatedAt: 200, | |||
| Properties: model.Properties{OriginModelName: "doubao-seedance-2-0-260128"}, | |||
| Data: []byte(`{"id":"cgt-upstream","status":"failed","error":{"code":"InvalidParameter","message":"duration is invalid"}}`), | |||
| } | |||
| data, err := buildAipingNativeFetchResponse(task) | |||
| require.NoError(t, err) | |||
| require.Contains(t, string(data), `"id":"task_public"`) | |||
| require.Contains(t, string(data), `"code":"InvalidParameter"`) | |||
| require.Contains(t, string(data), `"message":"duration is invalid"`) | |||
| require.NotContains(t, string(data), `"error":null`) | |||
| } | |||
| func newControllerJSONContext(t *testing.T, path string, body string) *gin.Context { | |||
| t.Helper() | |||
| w := httptest.NewRecorder() | |||
| c, _ := gin.CreateTestContext(w) | |||
| c.Request = httptest.NewRequest(http.MethodPost, path, strings.NewReader(body)) | |||
| c.Request.Header.Set("Content-Type", "application/json") | |||
| return c | |||
| } | |||
| @@ -0,0 +1,203 @@ | |||
| package controller | |||
| import ( | |||
| "context" | |||
| "fmt" | |||
| "io" | |||
| "net/http" | |||
| "net/url" | |||
| "strings" | |||
| "time" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/constant" | |||
| "github.com/QuantumNous/new-api/logger" | |||
| "github.com/QuantumNous/new-api/model" | |||
| "github.com/QuantumNous/new-api/service" | |||
| "github.com/QuantumNous/new-api/setting/system_setting" | |||
| "github.com/gin-gonic/gin" | |||
| ) | |||
| const defaultDoubaoAssetBaseURL = "https://ark.cn-beijing.volcengineapi.com" | |||
| var blockedAssetActions = map[string]struct{}{ | |||
| "createassetgroup": {}, | |||
| "getassetgroup": {}, | |||
| "listassetgroups": {}, | |||
| "updateassetgroup": {}, | |||
| "deleteassetgroup": {}, | |||
| } | |||
| func assetProxyError(c *gin.Context, status int, errType, message string) { | |||
| c.JSON(status, gin.H{ | |||
| "error": gin.H{ | |||
| "message": message, | |||
| "type": errType, | |||
| }, | |||
| }) | |||
| } | |||
| func buildDoubaoAssetURL(channel *model.Channel, action string, version string) (string, error) { | |||
| baseURL := defaultDoubaoAssetBaseURL | |||
| if channel != nil && channel.BaseURL != nil { | |||
| if configuredBaseURL := strings.TrimSpace(*channel.BaseURL); configuredBaseURL != "" { | |||
| baseURL = configuredBaseURL | |||
| } | |||
| } | |||
| u, err := url.Parse(baseURL) | |||
| if err != nil { | |||
| return "", err | |||
| } | |||
| if u.Scheme == "" || u.Host == "" { | |||
| return "", fmt.Errorf("invalid Doubao asset base URL: %s", baseURL) | |||
| } | |||
| u.Path = strings.TrimRight(u.Path, "/") + "/api/v1/multimodal/sd/assets" | |||
| u.RawQuery = "" | |||
| u.Fragment = "" | |||
| query := u.Query() | |||
| query.Set("Action", action) | |||
| if strings.TrimSpace(version) == "" { | |||
| version = "2024-01-01" | |||
| } | |||
| query.Set("Version", version) | |||
| u.RawQuery = query.Encode() | |||
| return u.String(), nil | |||
| } | |||
| func effectiveDoubaoAssetGroup(c *gin.Context) string { | |||
| if group := strings.TrimSpace(common.GetContextKeyString(c, constant.ContextKeyUsingGroup)); group != "" { | |||
| return group | |||
| } | |||
| return strings.TrimSpace(common.GetContextKeyString(c, constant.ContextKeyTokenGroup)) | |||
| } | |||
| func concreteDoubaoAssetGroupsForRequest(c *gin.Context, autoGroups func(string) []string) []string { | |||
| group := effectiveDoubaoAssetGroup(c) | |||
| if group != "auto" { | |||
| if group == "" { | |||
| return nil | |||
| } | |||
| return []string{group} | |||
| } | |||
| userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup) | |||
| groups := autoGroups(userGroup) | |||
| concreteGroups := make([]string, 0, len(groups)) | |||
| for _, candidate := range groups { | |||
| candidate = strings.TrimSpace(candidate) | |||
| if candidate == "" || candidate == "auto" { | |||
| continue | |||
| } | |||
| concreteGroups = append(concreteGroups, candidate) | |||
| } | |||
| return concreteGroups | |||
| } | |||
| func resolveDoubaoAssetChannelForRequest(c *gin.Context) (*model.Channel, string, error) { | |||
| userId := c.GetInt("id") | |||
| var lastErr error | |||
| for _, group := range concreteDoubaoAssetGroupsForRequest(c, service.GetUserAutoGroup) { | |||
| channel, err := service.ResolveDoubaoAssetChannel(userId, group) | |||
| if err != nil { | |||
| lastErr = err | |||
| logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to resolve Doubao asset channel for group %s: %v", group, err)) | |||
| continue | |||
| } | |||
| if channel != nil { | |||
| return channel, group, nil | |||
| } | |||
| } | |||
| if lastErr != nil { | |||
| return nil, "", lastErr | |||
| } | |||
| return nil, "", nil | |||
| } | |||
| func DoubaoAssetProxy(c *gin.Context) { | |||
| action := strings.TrimSpace(c.Query("Action")) | |||
| if action == "" { | |||
| assetProxyError(c, http.StatusBadRequest, "invalid_request_error", "Action query parameter is required") | |||
| return | |||
| } | |||
| if _, ok := blockedAssetActions[strings.ToLower(action)]; ok { | |||
| assetProxyError(c, http.StatusBadRequest, "invalid_request_error", fmt.Sprintf("Asset group API (%s) is not supported", action)) | |||
| return | |||
| } | |||
| version := strings.TrimSpace(c.Query("Version")) | |||
| if version == "" { | |||
| version = "2024-01-01" | |||
| } | |||
| channel, _, err := resolveDoubaoAssetChannelForRequest(c) | |||
| if err != nil { | |||
| assetProxyError(c, http.StatusBadGateway, "server_error", fmt.Sprintf("Failed to resolve Doubao asset channel: %v", err)) | |||
| return | |||
| } | |||
| if channel == nil { | |||
| assetProxyError(c, http.StatusBadGateway, "server_error", "Failed to resolve Doubao asset channel") | |||
| return | |||
| } | |||
| if strings.TrimSpace(channel.Key) == "" { | |||
| assetProxyError(c, http.StatusBadGateway, "server_error", "Doubao asset channel key is missing") | |||
| return | |||
| } | |||
| upstreamURL, err := buildDoubaoAssetURL(channel, action, version) | |||
| if err != nil { | |||
| assetProxyError(c, http.StatusBadGateway, "server_error", fmt.Sprintf("Failed to build upstream URL: %v", err)) | |||
| return | |||
| } | |||
| fetchSetting := system_setting.GetFetchSetting() | |||
| if err := common.ValidateURLWithFetchSetting(upstreamURL, fetchSetting.EnableSSRFProtection, fetchSetting.AllowPrivateIp, fetchSetting.DomainFilterMode, fetchSetting.IpFilterMode, fetchSetting.DomainList, fetchSetting.IpList, fetchSetting.AllowedPorts, fetchSetting.ApplyIPFilterForDomain); err != nil { | |||
| logger.LogError(c.Request.Context(), fmt.Sprintf("Doubao asset URL blocked: %v", err)) | |||
| assetProxyError(c, http.StatusForbidden, "server_error", fmt.Sprintf("request blocked: %v", err)) | |||
| return | |||
| } | |||
| client, err := service.GetHttpClientWithProxy(channel.GetSetting().Proxy) | |||
| if err != nil { | |||
| assetProxyError(c, http.StatusBadGateway, "server_error", fmt.Sprintf("Failed to create proxy client: %v", err)) | |||
| return | |||
| } | |||
| if client == nil { | |||
| assetProxyError(c, http.StatusBadGateway, "server_error", "Failed to create proxy client") | |||
| return | |||
| } | |||
| ctx, cancel := context.WithTimeout(c.Request.Context(), 60*time.Second) | |||
| defer cancel() | |||
| req, err := http.NewRequestWithContext(ctx, http.MethodPost, upstreamURL, c.Request.Body) | |||
| if err != nil { | |||
| assetProxyError(c, http.StatusBadGateway, "server_error", fmt.Sprintf("Failed to create upstream request: %v", err)) | |||
| return | |||
| } | |||
| req.Header.Set("Content-Type", "application/json") | |||
| req.Header.Set("Accept", "application/json") | |||
| req.Header.Set("Authorization", "Bearer "+strings.TrimSpace(channel.Key)) | |||
| resp, err := client.Do(req) | |||
| if err != nil { | |||
| logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to proxy Doubao asset request to %s: %s", upstreamURL, err.Error())) | |||
| assetProxyError(c, http.StatusBadGateway, "server_error", fmt.Sprintf("Failed to proxy Doubao asset request: %v", err)) | |||
| return | |||
| } | |||
| defer resp.Body.Close() | |||
| for key, values := range resp.Header { | |||
| for _, value := range values { | |||
| c.Writer.Header().Add(key, value) | |||
| } | |||
| } | |||
| c.Writer.WriteHeader(resp.StatusCode) | |||
| if _, err = io.Copy(c.Writer, resp.Body); err != nil { | |||
| logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to copy Doubao asset upstream response: %s", err.Error())) | |||
| } | |||
| } | |||
| @@ -0,0 +1,103 @@ | |||
| package controller | |||
| import ( | |||
| "net/http" | |||
| "net/http/httptest" | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/constant" | |||
| "github.com/QuantumNous/new-api/model" | |||
| "github.com/gin-gonic/gin" | |||
| "github.com/stretchr/testify/assert" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func setupDoubaoAssetProxyRouter(t *testing.T) *gin.Engine { | |||
| t.Helper() | |||
| oldMode := gin.Mode() | |||
| gin.SetMode(gin.TestMode) | |||
| t.Cleanup(func() { | |||
| gin.SetMode(oldMode) | |||
| }) | |||
| r := gin.New() | |||
| r.POST("/api/v1/volcengine/asset", DoubaoAssetProxy) | |||
| return r | |||
| } | |||
| func decodeDoubaoAssetErrorMessage(t *testing.T, body string) string { | |||
| t.Helper() | |||
| var payload struct { | |||
| Error struct { | |||
| Message string `json:"message"` | |||
| Type string `json:"type"` | |||
| } `json:"error"` | |||
| } | |||
| require.NoError(t, common.Unmarshal([]byte(body), &payload)) | |||
| return payload.Error.Message | |||
| } | |||
| func TestDoubaoAssetProxyMissingActionReturns400(t *testing.T) { | |||
| router := setupDoubaoAssetProxyRouter(t) | |||
| req := httptest.NewRequest(http.MethodPost, "/api/v1/volcengine/asset?Version=2024-01-01", nil) | |||
| w := httptest.NewRecorder() | |||
| router.ServeHTTP(w, req) | |||
| assert.Equal(t, http.StatusBadRequest, w.Code) | |||
| assert.Equal(t, "Action query parameter is required", decodeDoubaoAssetErrorMessage(t, w.Body.String())) | |||
| } | |||
| func TestDoubaoAssetProxyBlocksAssetGroupActionsCaseInsensitively(t *testing.T) { | |||
| router := setupDoubaoAssetProxyRouter(t) | |||
| for _, action := range []string{"CreateAssetGroup", "getassetgroup", "LISTASSETGROUPS", "UpdateAssetGroup", "deleteassetgroup"} { | |||
| t.Run(action, func(t *testing.T) { | |||
| req := httptest.NewRequest(http.MethodPost, "/api/v1/volcengine/asset?Action="+action, nil) | |||
| w := httptest.NewRecorder() | |||
| router.ServeHTTP(w, req) | |||
| assert.Equal(t, http.StatusBadRequest, w.Code) | |||
| assert.Contains(t, decodeDoubaoAssetErrorMessage(t, w.Body.String()), "Asset group API ("+action+") is not supported") | |||
| }) | |||
| } | |||
| } | |||
| func TestBuildDoubaoAssetURLDefaultBaseAndEscapedQuery(t *testing.T) { | |||
| got, err := buildDoubaoAssetURL(&model.Channel{}, "ApplyUploadInner&Space", "") | |||
| require.NoError(t, err) | |||
| assert.Equal(t, defaultDoubaoAssetBaseURL+"/api/v1/multimodal/sd/assets?Action=ApplyUploadInner%26Space&Version=2024-01-01", got) | |||
| } | |||
| func TestBuildDoubaoAssetURLExplicitBaseURLOverridesDefault(t *testing.T) { | |||
| baseURL := "https://example.com/custom/" | |||
| got, err := buildDoubaoAssetURL(&model.Channel{BaseURL: &baseURL}, "CommitUploadInner", "2025-02-03") | |||
| require.NoError(t, err) | |||
| assert.Equal(t, "https://example.com/custom/api/v1/multimodal/sd/assets?Action=CommitUploadInner&Version=2025-02-03", got) | |||
| } | |||
| func TestEffectiveDoubaoAssetGroupUsesUsingGroupBeforeBlankTokenGroup(t *testing.T) { | |||
| c, _ := gin.CreateTestContext(httptest.NewRecorder()) | |||
| common.SetContextKey(c, constant.ContextKeyTokenGroup, "") | |||
| common.SetContextKey(c, constant.ContextKeyUsingGroup, "default") | |||
| assert.Equal(t, "default", effectiveDoubaoAssetGroup(c)) | |||
| } | |||
| func TestConcreteDoubaoAssetGroupsForAutoUsesUserAutoGroups(t *testing.T) { | |||
| c, _ := gin.CreateTestContext(httptest.NewRecorder()) | |||
| common.SetContextKey(c, constant.ContextKeyUsingGroup, "auto") | |||
| common.SetContextKey(c, constant.ContextKeyTokenGroup, "") | |||
| groups := concreteDoubaoAssetGroupsForRequest(c, func(string) []string { | |||
| return []string{"default", "vip"} | |||
| }) | |||
| assert.Equal(t, []string{"default", "vip"}, groups) | |||
| } | |||
| @@ -0,0 +1,510 @@ | |||
| package controller | |||
| import ( | |||
| "bytes" | |||
| "fmt" | |||
| "io" | |||
| "net/http" | |||
| "strconv" | |||
| "strings" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/constant" | |||
| "github.com/QuantumNous/new-api/dto" | |||
| "github.com/QuantumNous/new-api/logger" | |||
| "github.com/QuantumNous/new-api/middleware" | |||
| "github.com/QuantumNous/new-api/model" | |||
| "github.com/QuantumNous/new-api/relay" | |||
| klingaiping "github.com/QuantumNous/new-api/relay/channel/task/kling/aiping" | |||
| relaycommon "github.com/QuantumNous/new-api/relay/common" | |||
| "github.com/QuantumNous/new-api/service" | |||
| "github.com/QuantumNous/new-api/types" | |||
| "github.com/gin-gonic/gin" | |||
| ) | |||
| func KlingAipingNativeTaskSubmit(c *gin.Context) { | |||
| route, ok := klingaiping.FindRoute(c.Request.Method, c.FullPath(), klingaiping.RouteKindSubmit) | |||
| if !ok { | |||
| c.JSON(http.StatusNotFound, gin.H{"code": 404, "message": "route not found"}) | |||
| return | |||
| } | |||
| relayInfo, err := relaycommon.GenRelayInfo(c, types.RelayFormatTask, nil, nil) | |||
| if err != nil { | |||
| c.JSON(http.StatusInternalServerError, gin.H{"code": 500, "message": err.Error()}) | |||
| return | |||
| } | |||
| payload, err := readJSONPayload(c) | |||
| if err != nil { | |||
| c.JSON(http.StatusBadRequest, gin.H{"code": 400, "message": err.Error()}) | |||
| return | |||
| } | |||
| configureKlingAipingTaskRelayInfo(c, relayInfo, route, payload) | |||
| if bound := resolveKlingAipingBoundChannel(c, relayInfo); bound != nil { | |||
| relayInfo.LockedChannel = bound | |||
| } | |||
| relayTaskWithInfo(c, relayInfo) | |||
| } | |||
| func resolveKlingAipingBoundChannel(c *gin.Context, relayInfo *relaycommon.RelayInfo) *model.Channel { | |||
| group := concreteTaskVideoBindingGroup(c, relayInfo.TokenGroup) | |||
| if group == "" { | |||
| return nil | |||
| } | |||
| channel, err := service.GetBoundKlingAssetChannelForModel(c.GetInt("id"), group, relayInfo.OriginModelName) | |||
| if err != nil { | |||
| logger.LogError(c, fmt.Sprintf("resolve kling aiping bound channel failed: %v", err)) | |||
| return nil | |||
| } | |||
| if channel == nil || !service.IsUsableKlingAssetChannel(channel, group) { | |||
| return nil | |||
| } | |||
| if !model.IsChannelEnabledForGroupModel(group, relayInfo.OriginModelName, channel.Id) { | |||
| return nil | |||
| } | |||
| return channel | |||
| } | |||
| func configureKlingAipingTaskRelayInfo(c *gin.Context, relayInfo *relaycommon.RelayInfo, route klingaiping.Route, payload map[string]any) { | |||
| // Native Kling routes bypass Distribute(), so force relayTaskWithInfo to | |||
| // select a channel instead of reading a preselected one from context. | |||
| if relayInfo.ChannelMeta == nil { | |||
| relayInfo.ChannelMeta = &relaycommon.ChannelMeta{} | |||
| } | |||
| modelName := resolveKlingAipingModel(payload, route) | |||
| relayInfo.OriginModelName = modelName | |||
| relayInfo.Action = route.Action | |||
| relaycommon.StoreTaskRequest(c, relayInfo, route.Action, relaycommon.TaskSubmitReq{ | |||
| Model: modelName, | |||
| Prompt: stringFromMap(payload, "prompt"), | |||
| Duration: durationFromMap(payload), | |||
| Metadata: payload, | |||
| }) | |||
| } | |||
| func durationFromMap(m map[string]any) int { | |||
| for _, key := range []string{"duration", "seconds"} { | |||
| if duration, ok := intFromMapValue(m[key]); ok { | |||
| return duration | |||
| } | |||
| } | |||
| return 0 | |||
| } | |||
| func intFromMapValue(value any) (int, bool) { | |||
| switch v := value.(type) { | |||
| case int: | |||
| return v, true | |||
| case int64: | |||
| return int(v), true | |||
| case float64: | |||
| return int(v), true | |||
| case string: | |||
| duration, err := strconv.Atoi(strings.TrimSpace(v)) | |||
| if err == nil { | |||
| return duration, true | |||
| } | |||
| } | |||
| return 0, false | |||
| } | |||
| func KlingAipingNativeTaskFetch(c *gin.Context) { | |||
| route, ok := klingaiping.FindRoute(c.Request.Method, c.FullPath(), klingaiping.RouteKindFetch) | |||
| if !ok { | |||
| c.JSON(http.StatusNotFound, gin.H{"code": 404, "message": "route not found"}) | |||
| return | |||
| } | |||
| task, ok := getKlingAipingUserTask(c, c.Param("task_id"), route.Action) | |||
| if !ok { | |||
| return | |||
| } | |||
| refreshKlingAipingTaskIfNeeded(task) | |||
| c.JSON(http.StatusOK, buildKlingAipingTaskPayload(task)) | |||
| } | |||
| func KlingAipingNativeTaskList(c *gin.Context) { | |||
| route, ok := klingaiping.FindRoute(c.Request.Method, c.FullPath(), klingaiping.RouteKindList) | |||
| if !ok { | |||
| c.JSON(http.StatusNotFound, gin.H{"code": 404, "message": "route not found"}) | |||
| return | |||
| } | |||
| pageNum, pageSize, err := parseKlingAipingPage(c) | |||
| if err != nil { | |||
| c.JSON(http.StatusUnprocessableEntity, gin.H{"code": 422, "message": err.Error()}) | |||
| return | |||
| } | |||
| tasks := model.TaskGetAllUserTask(c.GetInt("id"), (pageNum-1)*pageSize, pageSize, model.SyncTaskQueryParams{ | |||
| Platform: constant.TaskPlatform(strconv.Itoa(constant.ChannelTypeKlingAiping)), | |||
| Action: route.Action, | |||
| }) | |||
| data := make([]any, 0, len(tasks)) | |||
| for _, task := range tasks { | |||
| data = append(data, taskDataObject(task)) | |||
| } | |||
| c.JSON(http.StatusOK, gin.H{ | |||
| "code": 0, | |||
| "message": "success", | |||
| "request_id": c.GetString(common.RequestIdKey), | |||
| "data": data, | |||
| }) | |||
| } | |||
| func KlingAipingNativeProxy(c *gin.Context) { | |||
| route, ok := klingaiping.FindRoute(c.Request.Method, c.FullPath(), klingaiping.RouteKindProxy) | |||
| if !ok { | |||
| c.JSON(http.StatusNotFound, gin.H{"code": 404, "message": "route not found"}) | |||
| return | |||
| } | |||
| group := common.GetContextKeyString(c, constant.ContextKeyUsingGroup) | |||
| channel := resolveKlingAipingBoundChannelForProxy(c, group, route.BillingModel) | |||
| if channel == nil { | |||
| var err error | |||
| channel, _, err = service.CacheGetRandomSatisfiedChannel(&service.RetryParam{ | |||
| Ctx: c, | |||
| TokenGroup: group, | |||
| ModelName: route.BillingModel, | |||
| Retry: common.GetPointer(0), | |||
| AllowedChannelTypes: service.VideoAssetChannelTypesForFamily(service.VideoAssetFamilyKling), | |||
| }) | |||
| if err != nil { | |||
| c.JSON(http.StatusServiceUnavailable, gin.H{"code": 503, "message": err.Error()}) | |||
| return | |||
| } | |||
| } | |||
| if setupErr := middleware.SetupContextForSelectedChannel(c, channel, route.BillingModel); setupErr != nil { | |||
| c.JSON(setupErr.StatusCode, gin.H{"code": setupErr.GetErrorCode(), "message": setupErr.Error()}) | |||
| return | |||
| } | |||
| resp, err := doKlingAipingProxyRequest(c, route, channel) | |||
| if err != nil { | |||
| c.JSON(http.StatusBadGateway, gin.H{"code": 502, "message": err.Error()}) | |||
| return | |||
| } | |||
| defer resp.Body.Close() | |||
| copyProxyResponse(c, resp) | |||
| } | |||
| func resolveKlingAipingBoundChannelForProxy(c *gin.Context, group, billingModel string) *model.Channel { | |||
| group = strings.TrimSpace(group) | |||
| if group == "" || group == "auto" || strings.TrimSpace(billingModel) == "" { | |||
| return nil | |||
| } | |||
| channel, err := service.GetBoundKlingAssetChannelForModel(c.GetInt("id"), group, billingModel) | |||
| if err != nil { | |||
| logger.LogError(c, fmt.Sprintf("resolve kling aiping bound channel (proxy) failed: %v", err)) | |||
| return nil | |||
| } | |||
| if channel == nil || !service.IsUsableKlingAssetChannel(channel, group) { | |||
| return nil | |||
| } | |||
| if !model.IsChannelEnabledForGroupModel(group, billingModel, channel.Id) { | |||
| return nil | |||
| } | |||
| return channel | |||
| } | |||
| func readJSONPayload(c *gin.Context) (map[string]any, error) { | |||
| body, err := common.GetBodyStorage(c) | |||
| if err != nil { | |||
| return nil, err | |||
| } | |||
| data, err := body.Bytes() | |||
| if err != nil { | |||
| return nil, err | |||
| } | |||
| _, _ = body.Seek(0, io.SeekStart) | |||
| c.Request.Body = io.NopCloser(body) | |||
| payload := map[string]any{} | |||
| if strings.TrimSpace(string(data)) == "" { | |||
| return payload, nil | |||
| } | |||
| if err := common.Unmarshal(data, &payload); err != nil { | |||
| return nil, err | |||
| } | |||
| return payload, nil | |||
| } | |||
| func resolveKlingAipingModel(payload map[string]any, route klingaiping.Route) string { | |||
| if modelName := stringFromMap(payload, "model_name"); modelName != "" { | |||
| return modelName | |||
| } | |||
| if modelName := stringFromMap(payload, "model"); modelName != "" { | |||
| return modelName | |||
| } | |||
| if route.BillingModel != "" { | |||
| return route.BillingModel | |||
| } | |||
| return "kling-v3" | |||
| } | |||
| func stringFromMap(payload map[string]any, key string) string { | |||
| if value, ok := payload[key].(string); ok { | |||
| return strings.TrimSpace(value) | |||
| } | |||
| return "" | |||
| } | |||
| func getKlingAipingUserTask(c *gin.Context, taskID string, action string) (*model.Task, bool) { | |||
| task, exist, err := model.GetByTaskId(c.GetInt("id"), taskID) | |||
| if err != nil { | |||
| c.JSON(http.StatusInternalServerError, gin.H{"code": 500, "message": err.Error()}) | |||
| return nil, false | |||
| } | |||
| if !exist || task.Platform != constant.TaskPlatform(strconv.Itoa(constant.ChannelTypeKlingAiping)) || task.Action != action { | |||
| c.JSON(http.StatusBadRequest, gin.H{"code": 400, "message": "task not found"}) | |||
| return nil, false | |||
| } | |||
| return task, true | |||
| } | |||
| func refreshKlingAipingTaskIfNeeded(task *model.Task) { | |||
| if task.Status == model.TaskStatusSuccess || task.Status == model.TaskStatusFailure { | |||
| return | |||
| } | |||
| channelModel, err := model.GetChannelById(task.ChannelId, true) | |||
| if err != nil || channelModel == nil { | |||
| return | |||
| } | |||
| adaptor := relay.GetTaskAdaptor(task.Platform) | |||
| if adaptor == nil { | |||
| return | |||
| } | |||
| resp, err := adaptor.FetchTask(channelModel.GetBaseURL(), channelModel.Key, map[string]any{ | |||
| "task_id": task.GetUpstreamTaskID(), | |||
| "action": task.Action, | |||
| }, channelModel.GetSetting().Proxy) | |||
| if err != nil || resp == nil { | |||
| return | |||
| } | |||
| defer resp.Body.Close() | |||
| body, err := io.ReadAll(resp.Body) | |||
| if err != nil || resp.StatusCode < 200 || resp.StatusCode >= 300 { | |||
| return | |||
| } | |||
| taskInfo, err := adaptor.ParseTaskResult(body) | |||
| if err != nil || taskInfo == nil { | |||
| return | |||
| } | |||
| snap := task.Snapshot() | |||
| task.Data = body | |||
| if taskInfo.Status != "" { | |||
| task.Status = model.TaskStatus(taskInfo.Status) | |||
| } | |||
| if taskInfo.Progress != "" { | |||
| task.Progress = taskInfo.Progress | |||
| } | |||
| if taskInfo.Url != "" { | |||
| task.PrivateData.ResultURL = taskInfo.Url | |||
| } | |||
| if !snap.Equal(task.Snapshot()) { | |||
| _, _ = task.UpdateWithStatus(snap.Status) | |||
| } | |||
| } | |||
| func buildKlingAipingTaskPayload(task *model.Task) map[string]any { | |||
| payload := map[string]any{ | |||
| "code": 0, | |||
| "message": "success", | |||
| "data": taskDataObject(task), | |||
| } | |||
| return payload | |||
| } | |||
| func taskDataObject(task *model.Task) map[string]any { | |||
| payload := map[string]any{} | |||
| _ = common.Unmarshal(task.Data, &payload) | |||
| delete(payload, "aiping_id") | |||
| data, _ := payload["data"].(map[string]any) | |||
| if data == nil { | |||
| data = map[string]any{} | |||
| } | |||
| data["task_id"] = task.TaskID | |||
| if _, ok := data["task_status"]; !ok { | |||
| data["task_status"] = mapKlingAipingTaskStatus(task.Status) | |||
| } | |||
| if _, ok := data["task_status_msg"]; !ok { | |||
| data["task_status_msg"] = task.FailReason | |||
| } | |||
| if _, ok := data["created_at"]; !ok && task.CreatedAt != 0 { | |||
| data["created_at"] = task.CreatedAt | |||
| } | |||
| if _, ok := data["updated_at"]; !ok && task.UpdatedAt != 0 { | |||
| data["updated_at"] = task.UpdatedAt | |||
| } | |||
| ensureKlingAipingWatermarkURL(data) | |||
| return data | |||
| } | |||
| func mapKlingAipingTaskStatus(status model.TaskStatus) string { | |||
| switch status { | |||
| case model.TaskStatusSubmitted, model.TaskStatusQueued: | |||
| return "submitted" | |||
| case model.TaskStatusInProgress: | |||
| return "processing" | |||
| case model.TaskStatusSuccess: | |||
| return "succeed" | |||
| case model.TaskStatusFailure: | |||
| return "failed" | |||
| default: | |||
| return "processing" | |||
| } | |||
| } | |||
| func ensureKlingAipingWatermarkURL(data map[string]any) { | |||
| taskResult, _ := data["task_result"].(map[string]any) | |||
| if taskResult == nil { | |||
| return | |||
| } | |||
| videos, _ := taskResult["videos"].([]any) | |||
| for _, videoAny := range videos { | |||
| video, _ := videoAny.(map[string]any) | |||
| if video == nil { | |||
| continue | |||
| } | |||
| if _, ok := video["watermark_url"]; !ok { | |||
| video["watermark_url"] = "" | |||
| } | |||
| } | |||
| } | |||
| func parseKlingAipingPage(c *gin.Context) (int, int, error) { | |||
| pageNum := parseIntDefault(c.Query("pageNum"), 1) | |||
| pageSize := parseIntDefault(c.Query("pageSize"), 30) | |||
| if pageNum < 1 || pageNum > 1000 { | |||
| return 0, 0, fmt.Errorf("pageNum must be in [1, 1000]") | |||
| } | |||
| if pageSize < 1 || pageSize > 500 { | |||
| return 0, 0, fmt.Errorf("pageSize must be in [1, 500]") | |||
| } | |||
| return pageNum, pageSize, nil | |||
| } | |||
| func parseIntDefault(raw string, fallback int) int { | |||
| if strings.TrimSpace(raw) == "" { | |||
| return fallback | |||
| } | |||
| v, err := strconv.Atoi(raw) | |||
| if err != nil { | |||
| return -1 | |||
| } | |||
| return v | |||
| } | |||
| func doKlingAipingProxyRequest(c *gin.Context, route klingaiping.Route, channelModel *model.Channel) (*http.Response, error) { | |||
| baseURL := strings.TrimRight(channelModel.GetBaseURL(), "/") | |||
| upstreamPath := strings.Replace(route.UpstreamPath, ":id", c.Param("id"), 1) | |||
| url := baseURL + upstreamPath | |||
| if c.Request.URL.RawQuery != "" { | |||
| url += "?" + c.Request.URL.RawQuery | |||
| } | |||
| var body io.Reader | |||
| if c.Request.Method != http.MethodGet { | |||
| data, err := proxyBodyBytes(c, route) | |||
| if err != nil { | |||
| return nil, err | |||
| } | |||
| body = bytes.NewReader(data) | |||
| } | |||
| req, err := http.NewRequest(c.Request.Method, url, body) | |||
| if err != nil { | |||
| return nil, err | |||
| } | |||
| req.Header.Set("Accept", "application/json") | |||
| req.Header.Set("Content-Type", "application/json") | |||
| key := common.GetContextKeyString(c, constant.ContextKeyChannelKey) | |||
| if key == "" { | |||
| key = channelModel.Key | |||
| } | |||
| req.Header.Set("Authorization", "Bearer "+key) | |||
| client, err := service.GetHttpClientWithProxy(channelModel.GetSetting().Proxy) | |||
| if err != nil { | |||
| return nil, err | |||
| } | |||
| return client.Do(req) | |||
| } | |||
| func proxyBodyBytes(c *gin.Context, route klingaiping.Route) ([]byte, error) { | |||
| storage, err := common.GetBodyStorage(c) | |||
| if err != nil { | |||
| return nil, err | |||
| } | |||
| return storage.Bytes() | |||
| } | |||
| func copyProxyResponse(c *gin.Context, resp *http.Response) { | |||
| if resp.StatusCode >= http.StatusBadRequest { | |||
| copyNormalizedProxyError(c, resp) | |||
| return | |||
| } | |||
| for key, values := range resp.Header { | |||
| for _, value := range values { | |||
| c.Writer.Header().Add(key, value) | |||
| } | |||
| } | |||
| c.Status(resp.StatusCode) | |||
| _, _ = io.Copy(c.Writer, resp.Body) | |||
| } | |||
| func copyNormalizedProxyError(c *gin.Context, resp *http.Response) { | |||
| body, _ := io.ReadAll(resp.Body) | |||
| message := strings.TrimSpace(string(body)) | |||
| payload := map[string]any{} | |||
| if len(body) > 0 && common.Unmarshal(body, &payload) == nil { | |||
| if msg := stringFromMap(payload, "message"); msg != "" { | |||
| message = msg | |||
| } else if msg := stringFromMap(payload, "msg"); msg != "" { | |||
| message = msg | |||
| } else if detail, ok := payload["detail"]; ok { | |||
| if detailMap, ok := detail.(map[string]any); ok { | |||
| if msg := stringFromMap(detailMap, "message"); msg != "" { | |||
| message = msg | |||
| } else if msg := stringFromMap(detailMap, "msg"); msg != "" { | |||
| message = msg | |||
| } else { | |||
| message = fmt.Sprint(detail) | |||
| } | |||
| } else { | |||
| message = fmt.Sprint(detail) | |||
| } | |||
| } | |||
| delete(payload, "msg") | |||
| } else { | |||
| payload = map[string]any{} | |||
| } | |||
| if message == "" { | |||
| message = resp.Status | |||
| } | |||
| if _, ok := payload["code"]; !ok { | |||
| payload["code"] = resp.StatusCode | |||
| } | |||
| payload["message"] = message | |||
| payload["request_id"] = c.GetString(common.RequestIdKey) | |||
| c.JSON(resp.StatusCode, payload) | |||
| } | |||
| func normalizeKlingAipingTaskError(taskErr *dto.TaskError) { | |||
| if taskErr == nil || strings.TrimSpace(taskErr.Message) == "" { | |||
| return | |||
| } | |||
| payload := map[string]any{} | |||
| if common.Unmarshal([]byte(taskErr.Message), &payload) != nil { | |||
| return | |||
| } | |||
| if msg := stringFromMap(payload, "message"); msg != "" { | |||
| taskErr.Message = msg | |||
| return | |||
| } | |||
| if msg := stringFromMap(payload, "msg"); msg != "" { | |||
| taskErr.Message = msg | |||
| return | |||
| } | |||
| if detail, ok := payload["detail"].(map[string]any); ok { | |||
| if msg := stringFromMap(detail, "message"); msg != "" { | |||
| taskErr.Message = msg | |||
| } | |||
| } | |||
| } | |||
| @@ -0,0 +1,250 @@ | |||
| package controller | |||
| import ( | |||
| "io" | |||
| "net/http" | |||
| "net/http/httptest" | |||
| "strings" | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/constant" | |||
| "github.com/QuantumNous/new-api/dto" | |||
| "github.com/QuantumNous/new-api/model" | |||
| klingaiping "github.com/QuantumNous/new-api/relay/channel/task/kling/aiping" | |||
| relaycommon "github.com/QuantumNous/new-api/relay/common" | |||
| "github.com/QuantumNous/new-api/service" | |||
| "github.com/gin-gonic/gin" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestKlingAipingSubmitPreparationUsesModelNameBeforeModel(t *testing.T) { | |||
| c := newControllerJSONContext(t, "/v1/videos/text2video", `{ | |||
| "model":"Kling-V1.6", | |||
| "model_name":"Kling-V2.6", | |||
| "prompt":"prompt" | |||
| }`) | |||
| payload, err := readJSONPayload(c) | |||
| require.NoError(t, err) | |||
| route, ok := klingaiping.FindRoute(http.MethodPost, "/v1/videos/text2video", klingaiping.RouteKindSubmit) | |||
| require.True(t, ok) | |||
| info := &relaycommon.RelayInfo{TaskRelayInfo: &relaycommon.TaskRelayInfo{}} | |||
| modelName := resolveKlingAipingModel(payload, route) | |||
| relaycommon.StoreTaskRequest(c, info, route.Action, relaycommon.TaskSubmitReq{ | |||
| Model: modelName, | |||
| Prompt: stringFromMap(payload, "prompt"), | |||
| Metadata: payload, | |||
| }) | |||
| stored, err := relaycommon.GetTaskRequest(c) | |||
| require.NoError(t, err) | |||
| require.Equal(t, "Kling-V2.6", stored.Model) | |||
| require.Equal(t, "prompt", stored.Prompt) | |||
| require.Equal(t, "Kling-V1.6", stored.Metadata["model"]) | |||
| require.Equal(t, "Kling-V2.6", stored.Metadata["model_name"]) | |||
| } | |||
| func TestKlingAipingSubmitPreparationDoesNotLockChannel(t *testing.T) { | |||
| c := newControllerJSONContext(t, "/v1/general/custom-voices", `{ | |||
| "voice_url":"https://example.com/voice.mp3", | |||
| "voice_name":"voice" | |||
| }`) | |||
| payload, err := readJSONPayload(c) | |||
| require.NoError(t, err) | |||
| route, ok := klingaiping.FindRoute(http.MethodPost, "/v1/general/custom-voices", klingaiping.RouteKindSubmit) | |||
| require.True(t, ok) | |||
| info := &relaycommon.RelayInfo{TaskRelayInfo: &relaycommon.TaskRelayInfo{}} | |||
| info.OriginModelName = resolveKlingAipingModel(payload, route) | |||
| info.Action = route.Action | |||
| relaycommon.StoreTaskRequest(c, info, route.Action, relaycommon.TaskSubmitReq{ | |||
| Model: info.OriginModelName, | |||
| Metadata: payload, | |||
| }) | |||
| require.Equal(t, klingaiping.ModelCustomVoices, info.OriginModelName) | |||
| require.Equal(t, klingaiping.ActionVoicesCreate, info.Action) | |||
| require.Nil(t, info.LockedChannel) | |||
| } | |||
| func TestConfigureKlingAipingTaskRelayInfoForcesChannelSelection(t *testing.T) { | |||
| c := newControllerJSONContext(t, "/v1/videos/text2video", `{"model_name":"Kling-V2.6","prompt":"prompt"}`) | |||
| payload, err := readJSONPayload(c) | |||
| require.NoError(t, err) | |||
| route, ok := klingaiping.FindRoute(http.MethodPost, "/v1/videos/text2video", klingaiping.RouteKindSubmit) | |||
| require.True(t, ok) | |||
| info := &relaycommon.RelayInfo{TaskRelayInfo: &relaycommon.TaskRelayInfo{}} | |||
| configureKlingAipingTaskRelayInfo(c, info, route, payload) | |||
| require.NotNil(t, info.ChannelMeta) | |||
| require.Zero(t, info.ChannelMeta.ChannelType) | |||
| require.Equal(t, "Kling-V2.6", info.OriginModelName) | |||
| require.Equal(t, klingaiping.ActionText2Video, info.Action) | |||
| require.Nil(t, info.LockedChannel) | |||
| } | |||
| func TestConfigureKlingAipingTaskRelayInfoStoresDuration(t *testing.T) { | |||
| c := newControllerJSONContext(t, "/v1/videos/text2video", `{"model_name":"Kling-V2.6","prompt":"prompt","duration":5}`) | |||
| payload, err := readJSONPayload(c) | |||
| require.NoError(t, err) | |||
| route, ok := klingaiping.FindRoute(http.MethodPost, "/v1/videos/text2video", klingaiping.RouteKindSubmit) | |||
| require.True(t, ok) | |||
| info := &relaycommon.RelayInfo{TaskRelayInfo: &relaycommon.TaskRelayInfo{}} | |||
| configureKlingAipingTaskRelayInfo(c, info, route, payload) | |||
| stored, err := relaycommon.GetTaskRequest(c) | |||
| require.NoError(t, err) | |||
| require.Equal(t, 5, stored.Duration) | |||
| require.Empty(t, stored.Seconds) | |||
| } | |||
| func TestRequiredTaskChannelTypeOnlyMatchesKlingAipingNativePaths(t *testing.T) { | |||
| c := newControllerJSONContext(t, "/v1/videos/text2video", `{}`) | |||
| c.Request.URL.Path = "/v1/videos/text2video" | |||
| require.Zero(t, requiredTaskChannelTypeForRequest(c)) | |||
| require.ElementsMatch(t, | |||
| service.VideoAssetChannelTypesForFamily(service.VideoAssetFamilyKling), | |||
| allowedTaskChannelTypesForRequest(c), | |||
| ) | |||
| c = newControllerJSONContext(t, "/v1/videos/video-extend", `{}`) | |||
| c.Request.URL.Path = "/v1/videos/video-extend" | |||
| require.Zero(t, requiredTaskChannelTypeForRequest(c)) | |||
| require.ElementsMatch(t, | |||
| service.VideoAssetChannelTypesForFamily(service.VideoAssetFamilyKling), | |||
| allowedTaskChannelTypesForRequest(c), | |||
| ) | |||
| c = newControllerJSONContext(t, "/v1/videos/video_123/remix", `{}`) | |||
| c.Request.URL.Path = "/v1/videos/video_123/remix" | |||
| require.Zero(t, requiredTaskChannelTypeForRequest(c)) | |||
| require.Nil(t, allowedTaskChannelTypesForRequest(c)) | |||
| } | |||
| func TestKlingAipingTaskDataObjectUsesPublicTaskIDAndWatermarkURL(t *testing.T) { | |||
| task := &model.Task{ | |||
| TaskID: "task_public", | |||
| Status: model.TaskStatusSuccess, | |||
| CreatedAt: 100, | |||
| UpdatedAt: 200, | |||
| Data: []byte(`{ | |||
| "code":0, | |||
| "aiping_id":"internal", | |||
| "data":{ | |||
| "task_id":"899333358055493641", | |||
| "task_status":"succeed", | |||
| "task_result":{"videos":[{"id":"v1","url":"https://example.com/video.mp4","duration":"5.041"}]} | |||
| } | |||
| }`), | |||
| } | |||
| data := taskDataObject(task) | |||
| require.Equal(t, "task_public", data["task_id"]) | |||
| taskResult := data["task_result"].(map[string]any) | |||
| videos := taskResult["videos"].([]any) | |||
| require.Equal(t, "", videos[0].(map[string]any)["watermark_url"]) | |||
| } | |||
| func TestDoKlingAipingProxyRequestUsesSelectedContextKey(t *testing.T) { | |||
| service.InitHttpClient() | |||
| var gotAuth string | |||
| server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | |||
| gotAuth = r.Header.Get("Authorization") | |||
| w.Header().Set("Content-Type", "application/json") | |||
| _, _ = w.Write([]byte(`{"code":0,"message":"success","data":[]}`)) | |||
| })) | |||
| defer server.Close() | |||
| c := newControllerJSONContext(t, "/v1/general/advanced-presets-elements", ``) | |||
| common.SetContextKey(c, constant.ContextKeyChannelKey, "selected-key") | |||
| route, ok := klingaiping.FindRoute(http.MethodGet, "/v1/general/advanced-presets-elements", klingaiping.RouteKindProxy) | |||
| require.True(t, ok) | |||
| channel := &model.Channel{ | |||
| Key: "raw-channel-key", | |||
| BaseURL: common.GetPointer(server.URL), | |||
| } | |||
| resp, err := doKlingAipingProxyRequest(c, route, channel) | |||
| require.NoError(t, err) | |||
| defer resp.Body.Close() | |||
| require.Equal(t, "Bearer selected-key", gotAuth) | |||
| } | |||
| func TestCopyProxyResponseNormalizesMsgError(t *testing.T) { | |||
| w := httptest.NewRecorder() | |||
| c, _ := gin.CreateTestContext(w) | |||
| c.Set(common.RequestIdKey, "req-test") | |||
| resp := &http.Response{ | |||
| StatusCode: http.StatusUnauthorized, | |||
| Header: http.Header{"Content-Type": []string{"application/json"}}, | |||
| Body: io.NopCloser(strings.NewReader(`{"code":401,"msg":"unauthorized","data":null}`)), | |||
| } | |||
| copyProxyResponse(c, resp) | |||
| require.Equal(t, http.StatusUnauthorized, w.Code) | |||
| require.Contains(t, w.Body.String(), `"message":"unauthorized"`) | |||
| require.NotContains(t, w.Body.String(), `"msg"`) | |||
| require.Contains(t, w.Body.String(), `"request_id":"req-test"`) | |||
| } | |||
| func TestCopyProxyResponseNormalizesPlainTextError(t *testing.T) { | |||
| w := httptest.NewRecorder() | |||
| c, _ := gin.CreateTestContext(w) | |||
| c.Set(common.RequestIdKey, "req-test") | |||
| resp := &http.Response{ | |||
| StatusCode: http.StatusMethodNotAllowed, | |||
| Header: http.Header{"Content-Type": []string{"text/plain"}}, | |||
| Body: io.NopCloser(strings.NewReader("Method Not Allowed")), | |||
| } | |||
| copyProxyResponse(c, resp) | |||
| require.Equal(t, http.StatusMethodNotAllowed, w.Code) | |||
| require.Contains(t, w.Body.String(), `"message":"Method Not Allowed"`) | |||
| require.Contains(t, w.Body.String(), `"request_id":"req-test"`) | |||
| } | |||
| func TestCopyProxyResponseNormalizesDetailMessageError(t *testing.T) { | |||
| w := httptest.NewRecorder() | |||
| c, _ := gin.CreateTestContext(w) | |||
| c.Set(common.RequestIdKey, "req-test") | |||
| resp := &http.Response{ | |||
| StatusCode: http.StatusServiceUnavailable, | |||
| Header: http.Header{"Content-Type": []string{"application/json"}}, | |||
| Body: io.NopCloser(strings.NewReader(`{"detail":{"message":"not found","error_type":"not_found"},"aiping_id":"internal"}`)), | |||
| } | |||
| copyProxyResponse(c, resp) | |||
| require.Equal(t, http.StatusServiceUnavailable, w.Code) | |||
| require.Contains(t, w.Body.String(), `"message":"not found"`) | |||
| require.NotContains(t, w.Body.String(), `map[`) | |||
| require.Contains(t, w.Body.String(), `"request_id":"req-test"`) | |||
| } | |||
| func TestNormalizeKlingAipingTaskErrorMessageExtractsUpstreamJSONMessage(t *testing.T) { | |||
| taskErr := &dto.TaskError{ | |||
| Code: "fail_to_fetch_task", | |||
| Message: `{"code":400,"message":"ERROR: image download failed","request_id":"upstream"}`, | |||
| StatusCode: http.StatusBadRequest, | |||
| } | |||
| normalizeKlingAipingTaskError(taskErr) | |||
| require.Equal(t, "ERROR: image download failed", taskErr.Message) | |||
| } | |||
| func TestParseKlingAipingPageBoundaries(t *testing.T) { | |||
| c := newControllerJSONContext(t, "/v1/videos/text2video?pageNum=1001&pageSize=30", `{}`) | |||
| c.Request.URL.RawQuery = "pageNum=1001&pageSize=30" | |||
| _, _, err := parseKlingAipingPage(c) | |||
| require.ErrorContains(t, err, "pageNum") | |||
| c = newControllerJSONContext(t, "/v1/videos/text2video?pageNum=1&pageSize=501", `{}`) | |||
| c.Request.URL.RawQuery = "pageNum=1&pageSize=501" | |||
| _, _, err = parseKlingAipingPage(c) | |||
| require.ErrorContains(t, err, "pageSize") | |||
| } | |||
| @@ -0,0 +1,59 @@ | |||
| package controller | |||
| import ( | |||
| "net/http" | |||
| "strings" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/model" | |||
| "github.com/QuantumNous/new-api/setting/ratio_setting" | |||
| "github.com/QuantumNous/new-api/types" | |||
| "github.com/gin-gonic/gin" | |||
| ) | |||
| func GetModelDisplayPricingRules(c *gin.Context) { | |||
| common.ApiSuccess(c, ratio_setting.GetModelDisplayPricingCopy()) | |||
| } | |||
| func GetModelDisplayPricingRule(c *gin.Context) { | |||
| modelName := strings.TrimPrefix(c.Param("model"), "/") | |||
| items := ratio_setting.GetModelDisplayPricing(modelName) | |||
| if len(items) == 0 { | |||
| c.JSON(http.StatusNotFound, gin.H{"success": false, "message": "model display pricing not found"}) | |||
| return | |||
| } | |||
| common.ApiSuccess(c, items) | |||
| } | |||
| func UpdateModelDisplayPricingRule(c *gin.Context) { | |||
| modelName := strings.TrimPrefix(c.Param("model"), "/") | |||
| var items []types.ModelDisplayPricingItem | |||
| if err := common.DecodeJson(c.Request.Body, &items); err != nil { | |||
| c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": err.Error()}) | |||
| return | |||
| } | |||
| if err := ratio_setting.SetModelDisplayPricing(modelName, items); err != nil { | |||
| c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": err.Error()}) | |||
| return | |||
| } | |||
| if err := model.UpdateOption(ratio_setting.ModelDisplayPricingOptionKey, ratio_setting.ModelDisplayPricing2JSONString()); err != nil { | |||
| c.JSON(http.StatusInternalServerError, gin.H{"success": false, "message": err.Error()}) | |||
| return | |||
| } | |||
| model.RefreshPricing() | |||
| common.ApiSuccess(c, ratio_setting.GetModelDisplayPricing(modelName)) | |||
| } | |||
| func DeleteModelDisplayPricingRule(c *gin.Context) { | |||
| modelName := strings.TrimPrefix(c.Param("model"), "/") | |||
| if err := ratio_setting.DeleteModelDisplayPricing(modelName); err != nil { | |||
| c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": err.Error()}) | |||
| return | |||
| } | |||
| if err := model.UpdateOption(ratio_setting.ModelDisplayPricingOptionKey, ratio_setting.ModelDisplayPricing2JSONString()); err != nil { | |||
| c.JSON(http.StatusInternalServerError, gin.H{"success": false, "message": err.Error()}) | |||
| return | |||
| } | |||
| model.RefreshPricing() | |||
| common.ApiSuccess(c, nil) | |||
| } | |||
| @@ -0,0 +1,96 @@ | |||
| package controller | |||
| import ( | |||
| "net/http" | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/model" | |||
| "github.com/QuantumNous/new-api/setting/ratio_setting" | |||
| "github.com/QuantumNous/new-api/types" | |||
| "github.com/gin-gonic/gin" | |||
| "github.com/glebarez/sqlite" | |||
| "github.com/stretchr/testify/require" | |||
| "gorm.io/gorm" | |||
| ) | |||
| func setupModelDisplayPricingControllerTest(t *testing.T) *gin.Engine { | |||
| t.Helper() | |||
| gin.SetMode(gin.TestMode) | |||
| db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) | |||
| require.NoError(t, err) | |||
| require.NoError(t, db.AutoMigrate(&model.Option{})) | |||
| sqlDB, err := db.DB() | |||
| require.NoError(t, err) | |||
| sqlDB.SetMaxOpenConns(1) | |||
| origDB := model.DB | |||
| origUsingSQLite := common.UsingSQLite | |||
| origRedisEnabled := common.RedisEnabled | |||
| model.DB = db | |||
| common.UsingSQLite = true | |||
| common.RedisEnabled = false | |||
| model.InitOptionMap() | |||
| require.NoError(t, ratio_setting.UpdateModelDisplayPricingByJSONString(`{}`)) | |||
| t.Cleanup(func() { | |||
| model.DB = origDB | |||
| common.UsingSQLite = origUsingSQLite | |||
| common.RedisEnabled = origRedisEnabled | |||
| require.NoError(t, ratio_setting.UpdateModelDisplayPricingByJSONString(`{}`)) | |||
| require.NoError(t, sqlDB.Close()) | |||
| }) | |||
| r := gin.New() | |||
| r.GET("/api/option/model_display_pricing", GetModelDisplayPricingRules) | |||
| r.GET("/api/option/model_display_pricing/*model", GetModelDisplayPricingRule) | |||
| r.PUT("/api/option/model_display_pricing/*model", UpdateModelDisplayPricingRule) | |||
| r.DELETE("/api/option/model_display_pricing/*model", DeleteModelDisplayPricingRule) | |||
| return r | |||
| } | |||
| func TestModelDisplayPricingControllerCRUDAndWildcard(t *testing.T) { | |||
| r := setupModelDisplayPricingControllerTest(t) | |||
| items := []types.ModelDisplayPricingItem{ | |||
| {Specification: "768p-6s", OfficialSupplierTip: "768P 6s", Price: 0.33333, Unit: "second", SortOrder: 1}, | |||
| } | |||
| w := performJSONRequest(t, r, http.MethodPut, "/api/option/model_display_pricing/provider/kling-video", items) | |||
| require.Equal(t, http.StatusOK, w.Code) | |||
| require.Len(t, ratio_setting.GetModelDisplayPricing("provider/kling-video"), 1) | |||
| w = performJSONRequest(t, r, http.MethodGet, "/api/option/model_display_pricing/provider/kling-video", nil) | |||
| require.Equal(t, http.StatusOK, w.Code) | |||
| require.Contains(t, w.Body.String(), "768p-6s") | |||
| w = performJSONRequest(t, r, http.MethodGet, "/api/option/model_display_pricing", nil) | |||
| require.Equal(t, http.StatusOK, w.Code) | |||
| require.Contains(t, w.Body.String(), "provider/kling-video") | |||
| w = performJSONRequest(t, r, http.MethodDelete, "/api/option/model_display_pricing/provider/kling-video", nil) | |||
| require.Equal(t, http.StatusOK, w.Code) | |||
| require.Empty(t, ratio_setting.GetModelDisplayPricing("provider/kling-video")) | |||
| } | |||
| func TestModelDisplayPricingControllerValidationError(t *testing.T) { | |||
| r := setupModelDisplayPricingControllerTest(t) | |||
| items := []types.ModelDisplayPricingItem{{Specification: "", Price: 1, Unit: "second"}} | |||
| w := performJSONRequest(t, r, http.MethodPut, "/api/option/model_display_pricing/kling-video", items) | |||
| require.Equal(t, http.StatusBadRequest, w.Code) | |||
| require.Contains(t, w.Body.String(), "specification is required") | |||
| } | |||
| func TestModelDisplayPricingOptionMapHydratesCache(t *testing.T) { | |||
| r := setupModelDisplayPricingControllerTest(t) | |||
| require.NotNil(t, r) | |||
| require.NoError(t, model.UpdateOption(ratio_setting.ModelDisplayPricingOptionKey, `{ | |||
| "kling-video":[{"specification":"768p","price":1,"unit":"second"}] | |||
| }`)) | |||
| items := ratio_setting.GetModelDisplayPricing("kling-video") | |||
| require.Len(t, items, 1) | |||
| require.Equal(t, "768p", items[0].Specification) | |||
| } | |||
| @@ -97,6 +97,10 @@ func CreateModelMeta(c *gin.Context) { | |||
| return | |||
| } | |||
| if strings.TrimSpace(m.Endpoints) == "" { | |||
| m.Endpoints = `["openai"]` | |||
| } | |||
| if err := m.Insert(); err != nil { | |||
| common.ApiError(c, err) | |||
| return | |||
| @@ -135,6 +139,10 @@ func UpdateModelMeta(c *gin.Context) { | |||
| return | |||
| } | |||
| if strings.TrimSpace(m.Endpoints) == "" { | |||
| m.Endpoints = `["openai"]` | |||
| } | |||
| if err := m.Update(); err != nil { | |||
| common.ApiError(c, err) | |||
| return | |||
| @@ -0,0 +1,89 @@ | |||
| package controller | |||
| import ( | |||
| "bytes" | |||
| "net/http" | |||
| "net/http/httptest" | |||
| "strconv" | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/model" | |||
| "github.com/gin-gonic/gin" | |||
| "github.com/glebarez/sqlite" | |||
| "github.com/stretchr/testify/require" | |||
| "gorm.io/gorm" | |||
| ) | |||
| func setupModelMetaTestDB(t *testing.T) { | |||
| t.Helper() | |||
| db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) | |||
| require.NoError(t, err) | |||
| require.NoError(t, db.AutoMigrate(&model.Model{})) | |||
| origDB := model.DB | |||
| model.DB = db | |||
| t.Cleanup(func() { | |||
| model.RefreshPricing() | |||
| model.DB = origDB | |||
| }) | |||
| } | |||
| func TestCreateModelMetaDefaultsEmptyEndpointsToOpenAI(t *testing.T) { | |||
| setupModelMetaTestDB(t) | |||
| gin.SetMode(gin.TestMode) | |||
| r := gin.New() | |||
| r.POST("/api/models/", CreateModelMeta) | |||
| req := httptest.NewRequest( | |||
| http.MethodPost, | |||
| "/api/models/", | |||
| bytes.NewBufferString(`{"model_name":"custom-chat-model"}`), | |||
| ) | |||
| req.Header.Set("Content-Type", "application/json") | |||
| w := httptest.NewRecorder() | |||
| r.ServeHTTP(w, req) | |||
| require.Equal(t, http.StatusOK, w.Code) | |||
| var resp map[string]any | |||
| require.NoError(t, common.Unmarshal(w.Body.Bytes(), &resp)) | |||
| require.True(t, resp["success"].(bool)) | |||
| var saved model.Model | |||
| require.NoError(t, model.DB.Where("model_name = ?", "custom-chat-model").First(&saved).Error) | |||
| require.Equal(t, `["openai"]`, saved.Endpoints) | |||
| } | |||
| func TestUpdateModelMetaDefaultsEmptyEndpointsToOpenAI(t *testing.T) { | |||
| setupModelMetaTestDB(t) | |||
| gin.SetMode(gin.TestMode) | |||
| existing := &model.Model{ | |||
| ModelName: "custom-chat-model", | |||
| Endpoints: `{"anthropic":{"path":"/v1/messages","method":"POST"}}`, | |||
| Status: 1, | |||
| } | |||
| require.NoError(t, existing.Insert()) | |||
| r := gin.New() | |||
| r.PUT("/api/models/", UpdateModelMeta) | |||
| req := httptest.NewRequest( | |||
| http.MethodPut, | |||
| "/api/models/", | |||
| bytes.NewBufferString(`{"id":`+strconv.Itoa(existing.Id)+`,"model_name":"custom-chat-model","endpoints":"","status":1}`), | |||
| ) | |||
| req.Header.Set("Content-Type", "application/json") | |||
| w := httptest.NewRecorder() | |||
| r.ServeHTTP(w, req) | |||
| require.Equal(t, http.StatusOK, w.Code) | |||
| var resp map[string]any | |||
| require.NoError(t, common.Unmarshal(w.Body.Bytes(), &resp)) | |||
| require.True(t, resp["success"].(bool)) | |||
| var saved model.Model | |||
| require.NoError(t, model.DB.First(&saved, existing.Id).Error) | |||
| require.Equal(t, `["openai"]`, saved.Endpoints) | |||
| } | |||
| @@ -0,0 +1,59 @@ | |||
| package controller | |||
| import ( | |||
| "net/http" | |||
| "strings" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/model" | |||
| "github.com/QuantumNous/new-api/setting/ratio_setting" | |||
| "github.com/QuantumNous/new-api/types" | |||
| "github.com/gin-gonic/gin" | |||
| ) | |||
| func GetModelPricingRules(c *gin.Context) { | |||
| common.ApiSuccess(c, ratio_setting.GetPricingConfigCopy()) | |||
| } | |||
| func GetModelPricingRule(c *gin.Context) { | |||
| modelName := strings.TrimPrefix(c.Param("model"), "/") | |||
| cfg := ratio_setting.GetPricingConfig(modelName) | |||
| if cfg == nil { | |||
| c.JSON(http.StatusNotFound, gin.H{"success": false, "message": "model pricing rule not found"}) | |||
| return | |||
| } | |||
| common.ApiSuccess(c, cfg) | |||
| } | |||
| func UpdateModelPricingRule(c *gin.Context) { | |||
| modelName := strings.TrimPrefix(c.Param("model"), "/") | |||
| var cfg types.PricingConfig | |||
| if err := common.DecodeJson(c.Request.Body, &cfg); err != nil { | |||
| c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": err.Error()}) | |||
| return | |||
| } | |||
| if err := ratio_setting.SetPricingConfig(modelName, &cfg); err != nil { | |||
| c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": err.Error()}) | |||
| return | |||
| } | |||
| if err := model.UpdateOption(ratio_setting.ModelPricingRulesOptionKey, ratio_setting.ModelPricingRules2JSONString()); err != nil { | |||
| c.JSON(http.StatusInternalServerError, gin.H{"success": false, "message": err.Error()}) | |||
| return | |||
| } | |||
| model.RefreshPricing() | |||
| common.ApiSuccess(c, ratio_setting.GetPricingConfig(modelName)) | |||
| } | |||
| func DeleteModelPricingRule(c *gin.Context) { | |||
| modelName := strings.TrimPrefix(c.Param("model"), "/") | |||
| if err := ratio_setting.DeletePricingConfig(modelName); err != nil { | |||
| c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": err.Error()}) | |||
| return | |||
| } | |||
| if err := model.UpdateOption(ratio_setting.ModelPricingRulesOptionKey, ratio_setting.ModelPricingRules2JSONString()); err != nil { | |||
| c.JSON(http.StatusInternalServerError, gin.H{"success": false, "message": err.Error()}) | |||
| return | |||
| } | |||
| model.RefreshPricing() | |||
| common.ApiSuccess(c, nil) | |||
| } | |||
| @@ -0,0 +1,153 @@ | |||
| package controller | |||
| import ( | |||
| "bytes" | |||
| "net/http" | |||
| "net/http/httptest" | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/model" | |||
| "github.com/QuantumNous/new-api/setting/ratio_setting" | |||
| "github.com/QuantumNous/new-api/types" | |||
| "github.com/gin-gonic/gin" | |||
| "github.com/glebarez/sqlite" | |||
| "github.com/stretchr/testify/assert" | |||
| "github.com/stretchr/testify/require" | |||
| "gorm.io/gorm" | |||
| ) | |||
| func setupModelPricingControllerTest(t *testing.T) *gin.Engine { | |||
| t.Helper() | |||
| gin.SetMode(gin.TestMode) | |||
| db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) | |||
| require.NoError(t, err) | |||
| require.NoError(t, db.AutoMigrate(&model.Option{})) | |||
| sqlDB, err := db.DB() | |||
| require.NoError(t, err) | |||
| sqlDB.SetMaxOpenConns(1) | |||
| origDB := model.DB | |||
| origUsingSQLite := common.UsingSQLite | |||
| origRedisEnabled := common.RedisEnabled | |||
| model.DB = db | |||
| common.UsingSQLite = true | |||
| common.RedisEnabled = false | |||
| model.InitOptionMap() | |||
| require.NoError(t, ratio_setting.UpdateModelPricingRulesByJSONString(`{}`)) | |||
| require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{"video-test":0.25}`)) | |||
| t.Cleanup(func() { | |||
| model.DB = origDB | |||
| common.UsingSQLite = origUsingSQLite | |||
| common.RedisEnabled = origRedisEnabled | |||
| require.NoError(t, ratio_setting.UpdateModelPricingRulesByJSONString(`{}`)) | |||
| require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{}`)) | |||
| require.NoError(t, sqlDB.Close()) | |||
| }) | |||
| r := gin.New() | |||
| r.GET("/api/option/model_pricing", GetModelPricingRules) | |||
| r.GET("/api/option/model_pricing/*model", GetModelPricingRule) | |||
| r.PUT("/api/option/model_pricing/*model", UpdateModelPricingRule) | |||
| r.DELETE("/api/option/model_pricing/*model", DeleteModelPricingRule) | |||
| return r | |||
| } | |||
| func TestModelPricingController_CRUD(t *testing.T) { | |||
| r := setupModelPricingControllerTest(t) | |||
| cfg := validControllerPricingConfig() | |||
| w := performJSONRequest(t, r, http.MethodPut, "/api/option/model_pricing/video-test", cfg) | |||
| require.Equal(t, http.StatusOK, w.Code) | |||
| assert.True(t, responseSuccess(t, w.Body.Bytes())) | |||
| w = performJSONRequest(t, r, http.MethodGet, "/api/option/model_pricing/video-test", nil) | |||
| require.Equal(t, http.StatusOK, w.Code) | |||
| assert.True(t, responseSuccess(t, w.Body.Bytes())) | |||
| assert.Contains(t, w.Body.String(), `"billing_unit":"per_call"`) | |||
| w = performJSONRequest(t, r, http.MethodGet, "/api/option/model_pricing", nil) | |||
| require.Equal(t, http.StatusOK, w.Code) | |||
| assert.True(t, responseSuccess(t, w.Body.Bytes())) | |||
| assert.Contains(t, w.Body.String(), "video-test") | |||
| w = performJSONRequest(t, r, http.MethodDelete, "/api/option/model_pricing/video-test", nil) | |||
| require.Equal(t, http.StatusOK, w.Code) | |||
| assert.True(t, responseSuccess(t, w.Body.Bytes())) | |||
| w = performJSONRequest(t, r, http.MethodGet, "/api/option/model_pricing/video-test", nil) | |||
| assert.Equal(t, http.StatusNotFound, w.Code) | |||
| legacyPrice, ok := ratio_setting.GetModelPrice("video-test", false) | |||
| require.True(t, ok) | |||
| assert.Equal(t, 0.25, legacyPrice) | |||
| } | |||
| func TestModelPricingController_PutRejectsBadConfig(t *testing.T) { | |||
| r := setupModelPricingControllerTest(t) | |||
| cfg := validControllerPricingConfig() | |||
| cfg.Table = nil | |||
| w := performJSONRequest(t, r, http.MethodPut, "/api/option/model_pricing/video-test", cfg) | |||
| assert.Equal(t, http.StatusBadRequest, w.Code) | |||
| assert.Contains(t, w.Body.String(), "table must not be empty") | |||
| } | |||
| func TestModelPricingController_CRUDWithSlashModelName(t *testing.T) { | |||
| r := setupModelPricingControllerTest(t) | |||
| cfg := validControllerPricingConfig() | |||
| modelName := "openai/gpt-4o" | |||
| w := performJSONRequest(t, r, http.MethodPut, "/api/option/model_pricing/"+modelName, cfg) | |||
| require.Equal(t, http.StatusOK, w.Code) | |||
| require.NotNil(t, ratio_setting.GetPricingConfig(modelName)) | |||
| w = performJSONRequest(t, r, http.MethodGet, "/api/option/model_pricing/"+modelName, nil) | |||
| require.Equal(t, http.StatusOK, w.Code) | |||
| assert.Contains(t, w.Body.String(), `"billing_unit":"per_call"`) | |||
| w = performJSONRequest(t, r, http.MethodDelete, "/api/option/model_pricing/"+modelName, nil) | |||
| require.Equal(t, http.StatusOK, w.Code) | |||
| assert.Nil(t, ratio_setting.GetPricingConfig(modelName)) | |||
| } | |||
| func performJSONRequest(t *testing.T, r *gin.Engine, method, path string, body any) *httptest.ResponseRecorder { | |||
| t.Helper() | |||
| var payload []byte | |||
| if body != nil { | |||
| var err error | |||
| payload, err = common.Marshal(body) | |||
| require.NoError(t, err) | |||
| } | |||
| req := httptest.NewRequest(method, path, bytes.NewReader(payload)) | |||
| req.Header.Set("Content-Type", "application/json") | |||
| w := httptest.NewRecorder() | |||
| r.ServeHTTP(w, req) | |||
| return w | |||
| } | |||
| func responseSuccess(t *testing.T, payload []byte) bool { | |||
| t.Helper() | |||
| var body struct { | |||
| Success bool `json:"success"` | |||
| } | |||
| require.NoError(t, common.Unmarshal(payload, &body)) | |||
| return body.Success | |||
| } | |||
| func validControllerPricingConfig() types.PricingConfig { | |||
| return types.PricingConfig{ | |||
| SchemaVersion: 1, | |||
| Scope: types.PricingScopeModel, | |||
| BillingUnit: types.BillingUnitPerCall, | |||
| PreconsumeStrategy: types.PreconsumeStrategyExact, | |||
| Dimensions: []types.PricingDimension{ | |||
| {Key: "resolution", Source: "request.resolution", Type: "string"}, | |||
| }, | |||
| Table: []types.PricingRow{ | |||
| {"resolution": "720P", "price": 0.1, "source": types.PricingRowSourceManual}, | |||
| }, | |||
| Fallback: types.PricingFallback{Strategy: types.PricingFallbackReject}, | |||
| } | |||
| } | |||
| @@ -3,12 +3,14 @@ package controller | |||
| import ( | |||
| "fmt" | |||
| "math" | |||
| "sort" | |||
| "strings" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/model" | |||
| "github.com/QuantumNous/new-api/service" | |||
| "github.com/QuantumNous/new-api/setting/ratio_setting" | |||
| "github.com/QuantumNous/new-api/types" | |||
| "github.com/gin-gonic/gin" | |||
| ) | |||
| @@ -121,21 +123,31 @@ func GetUserPricing(c *gin.Context) { | |||
| return | |||
| } | |||
| totalRatio := groupRatio | |||
| savingsPercent := int(math.Round((1 - totalRatio) * 100)) | |||
| result := gin.H{ | |||
| "success": true, | |||
| "model_name": pricingData.ModelName, | |||
| "quota_type": pricingData.QuotaType, | |||
| "group": userGroup, | |||
| "group_ratio": groupRatio, | |||
| "savings_percent": savingsPercent, | |||
| "logged_in": true, | |||
| "success": true, | |||
| "model_name": pricingData.ModelName, | |||
| "quota_type": pricingData.QuotaType, | |||
| "group": userGroup, | |||
| "group_ratio": groupRatio, | |||
| "savings_percent": savingsPercent, | |||
| "logged_in": true, | |||
| } | |||
| if pricingData.PricingConfig != nil { | |||
| result["pricing_config"] = pricingData.PricingConfig | |||
| } | |||
| if len(pricingData.DisplayPricing) > 0 { | |||
| result["display_pricing"] = pricingData.DisplayPricing | |||
| } | |||
| if pricingData.QuotaType == model.QuotaTypeByTokens { | |||
| if pricingData.PricingConfig != nil { | |||
| originalPrice := representativeMatrixPrice(pricingData.PricingConfig, pricingData.ModelPrice) | |||
| result["original_price"] = originalPrice | |||
| result["user_price"] = originalPrice * totalRatio | |||
| } | |||
| originalInput := pricingData.ModelRatio * 2 | |||
| originalOutput := pricingData.ModelRatio * pricingData.CompletionRatio * 2 | |||
| result["original_input"] = originalInput | |||
| @@ -163,14 +175,63 @@ func respondOriginalPrice(c *gin.Context, pricingData *model.Pricing) { | |||
| "logged_in": false, | |||
| } | |||
| if pricingData.QuotaType == model.QuotaTypeByTokens { | |||
| if pricingData.PricingConfig != nil { | |||
| result["pricing_config"] = pricingData.PricingConfig | |||
| result["original_price"] = representativeMatrixPrice(pricingData.PricingConfig, pricingData.ModelPrice) | |||
| } | |||
| result["original_input"] = pricingData.ModelRatio * 2 | |||
| result["original_output"] = pricingData.ModelRatio * pricingData.CompletionRatio * 2 | |||
| } else { | |||
| if pricingData.PricingConfig != nil { | |||
| result["pricing_config"] = pricingData.PricingConfig | |||
| } | |||
| result["original_price"] = pricingData.ModelPrice | |||
| } | |||
| if len(pricingData.DisplayPricing) > 0 { | |||
| result["display_pricing"] = pricingData.DisplayPricing | |||
| } | |||
| c.JSON(200, result) | |||
| } | |||
| func representativeMatrixPrice(cfg *types.PricingConfig, fallback float64) float64 { | |||
| if cfg == nil || len(cfg.Table) == 0 { | |||
| return fallback | |||
| } | |||
| prices := make([]float64, 0, len(cfg.Table)) | |||
| for _, row := range cfg.Table { | |||
| price, ok := pricingRowFloat(row["price"]) | |||
| if ok { | |||
| prices = append(prices, price) | |||
| } | |||
| } | |||
| if len(prices) == 0 { | |||
| return fallback | |||
| } | |||
| sort.Float64s(prices) | |||
| mid := len(prices) / 2 | |||
| if len(prices)%2 == 1 { | |||
| return prices[mid] | |||
| } | |||
| return (prices[mid-1] + prices[mid]) / 2 | |||
| } | |||
| func pricingRowFloat(value any) (float64, bool) { | |||
| switch n := value.(type) { | |||
| case float64: | |||
| return n, true | |||
| case float32: | |||
| return float64(n), true | |||
| case int: | |||
| return float64(n), true | |||
| case int64: | |||
| return float64(n), true | |||
| case int32: | |||
| return float64(n), true | |||
| default: | |||
| return 0, false | |||
| } | |||
| } | |||
| // formatDiscount 将倍率转换为中文折扣格式 | |||
| func formatDiscount(ratio float64) string { | |||
| if ratio <= 0 { | |||
| @@ -9,8 +9,9 @@ import ( | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/model" | |||
| "github.com/QuantumNous/new-api/setting/ratio_setting" | |||
| "github.com/glebarez/sqlite" | |||
| "github.com/QuantumNous/new-api/types" | |||
| "github.com/gin-gonic/gin" | |||
| "github.com/glebarez/sqlite" | |||
| "github.com/stretchr/testify/assert" | |||
| "github.com/stretchr/testify/require" | |||
| "gorm.io/gorm" | |||
| @@ -152,6 +153,74 @@ func TestGetUserPricing_NotLoggedIn_PerCall(t *testing.T) { | |||
| assert.Equal(t, float64(0.04), resp["original_price"]) | |||
| } | |||
| func TestGetUserPricingIncludesDisplayPricingWithoutChangingComputedPrice(t *testing.T) { | |||
| setupPricingTestDB(t) | |||
| modelName := "display-pricing-user-model" | |||
| router := gin.New() | |||
| router.GET("/api/pricing/user/*model", GetUserPricing) | |||
| setPricingCache([]model.Pricing{ | |||
| { | |||
| ModelName: modelName, | |||
| QuotaType: model.QuotaTypeByCall, | |||
| ModelPrice: 0.5, | |||
| EnableGroup: []string{"default"}, | |||
| DisplayPricing: []types.ModelDisplayPricingItem{ | |||
| {Specification: "768p-6s", Price: 0.33333, Unit: "second", SortOrder: 1, DiscountRate: 1}, | |||
| }, | |||
| }, | |||
| }) | |||
| w := httptest.NewRecorder() | |||
| req, _ := http.NewRequest("GET", "/api/pricing/user/"+modelName, nil) | |||
| router.ServeHTTP(w, req) | |||
| var resp map[string]interface{} | |||
| require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) | |||
| require.Equal(t, http.StatusOK, w.Code) | |||
| require.Contains(t, resp, "display_pricing") | |||
| require.Equal(t, float64(0.5), resp["original_price"]) | |||
| require.NotEqual(t, float64(0.33333), resp["original_price"]) | |||
| } | |||
| func TestGetPricingIncludesDisplayPricingForModelSquare(t *testing.T) { | |||
| setupPricingTestDB(t) | |||
| modelName := "display-pricing-square-model" | |||
| router := gin.New() | |||
| router.GET("/api/pricing", GetPricing) | |||
| setPricingCache([]model.Pricing{ | |||
| { | |||
| ModelName: modelName, | |||
| QuotaType: model.QuotaTypeByCall, | |||
| ModelPrice: 0.5, | |||
| EnableGroup: []string{"default"}, | |||
| DisplayPricing: []types.ModelDisplayPricingItem{ | |||
| {Specification: "768p-6s", Price: 0.33333, Unit: "second", SortOrder: 1, DiscountRate: 1}, | |||
| }, | |||
| }, | |||
| }) | |||
| withGroupRatio(t, `{"default":1}`) | |||
| w := httptest.NewRecorder() | |||
| req, _ := http.NewRequest("GET", "/api/pricing", nil) | |||
| router.ServeHTTP(w, req) | |||
| var resp map[string]interface{} | |||
| require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) | |||
| require.Equal(t, http.StatusOK, w.Code) | |||
| require.Equal(t, true, resp["success"]) | |||
| data := resp["data"].([]interface{}) | |||
| require.Len(t, data, 1) | |||
| item := data[0].(map[string]interface{}) | |||
| require.Equal(t, modelName, item["model_name"]) | |||
| require.Contains(t, item, "display_pricing") | |||
| displayPricing := item["display_pricing"].([]interface{}) | |||
| require.Len(t, displayPricing, 1) | |||
| require.Equal(t, "768p-6s", displayPricing[0].(map[string]interface{})["specification"]) | |||
| } | |||
| // TestGetUserPricing_LoggedIn_NoDiscount 已登录但无折扣(分组倍率=1,无个人倍率) | |||
| func TestGetUserPricing_LoggedIn_NoDiscount(t *testing.T) { | |||
| db := setupPricingTestDB(t) | |||
| @@ -226,8 +295,6 @@ func TestGetUserPricing_LoggedIn_GroupDiscount(t *testing.T) { | |||
| assert.Equal(t, "8折", resp["discount"]) | |||
| } | |||
| // TestGetUserPricing_PerCall_WithDiscount 按次计费 + 折扣 | |||
| func TestGetUserPricing_PerCall_WithDiscount(t *testing.T) { | |||
| db := setupPricingTestDB(t) | |||
| @@ -280,4 +347,3 @@ func TestFormatDiscount(t *testing.T) { | |||
| assert.Equal(t, tt.expected, result, "ratio=%.2f", tt.ratio) | |||
| } | |||
| } | |||
| @@ -24,6 +24,7 @@ import ( | |||
| "github.com/QuantumNous/new-api/service" | |||
| "github.com/QuantumNous/new-api/setting" | |||
| "github.com/QuantumNous/new-api/setting/operation_setting" | |||
| "github.com/QuantumNous/new-api/setting/ratio_setting" | |||
| "github.com/QuantumNous/new-api/types" | |||
| "github.com/bytedance/gopkg/util/gopool" | |||
| @@ -190,21 +191,21 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { | |||
| }() | |||
| retryParam := &service.RetryParam{ | |||
| Ctx: c, | |||
| TokenGroup: relayInfo.TokenGroup, | |||
| ModelName: relayInfo.OriginModelName, | |||
| Retry: common.GetPointer(0), | |||
| Ctx: c, | |||
| TokenGroup: relayInfo.TokenGroup, | |||
| ModelName: relayInfo.OriginModelName, | |||
| Retry: common.GetPointer(0), | |||
| RequireMatrixUsageBilling: relayInfo.RequireMatrixUsageBilling, | |||
| } | |||
| for ; retryParam.GetRetry() <= common.RetryTimes; retryParam.IncreaseRetry() { | |||
| channel, channelErr := getChannel(c, relayInfo, retryParam) | |||
| channel, _, channelErr := getChannel(c, relayInfo, retryParam) | |||
| if channelErr != nil { | |||
| logger.LogError(c, channelErr.Error()) | |||
| newAPIError = channelErr | |||
| break | |||
| } | |||
| addUsedChannel(c, channel.Id) | |||
| bodyStorage, bodyErr := common.GetBodyStorage(c) | |||
| if bodyErr != nil { | |||
| @@ -295,7 +296,37 @@ func fastTokenCountMetaForPricing(request dto.Request) *types.TokenCountMeta { | |||
| return meta | |||
| } | |||
| func getChannel(c *gin.Context, info *relaycommon.RelayInfo, retryParam *service.RetryParam) (*model.Channel, *types.NewAPIError) { | |||
| func concreteTaskVideoBindingGroup(c *gin.Context, group string) string { | |||
| group = strings.TrimSpace(group) | |||
| if group == "" { | |||
| group = strings.TrimSpace(common.GetContextKeyString(c, constant.ContextKeyTokenGroup)) | |||
| } | |||
| if group == "auto" { | |||
| autoGroup := strings.TrimSpace(common.GetContextKeyString(c, constant.ContextKeyAutoGroup)) | |||
| if autoGroup == "" || autoGroup == "auto" { | |||
| return "" | |||
| } | |||
| return autoGroup | |||
| } | |||
| return group | |||
| } | |||
| func persistTaskVideoBindingIfNeeded(userId int, group string, channel *model.Channel) error { | |||
| group = strings.TrimSpace(group) | |||
| if group == "" || group == "auto" || channel == nil { | |||
| return nil | |||
| } | |||
| family, ok := service.VideoAssetFamilyForChannelType(channel.Type) | |||
| if !ok { | |||
| return nil | |||
| } | |||
| if !service.IsUsableVideoAssetChannelForFamily(channel, group, "", family) { | |||
| return nil | |||
| } | |||
| return service.BindVideoAssetChannel(userId, group, channel, family) | |||
| } | |||
| func getChannel(c *gin.Context, info *relaycommon.RelayInfo, retryParam *service.RetryParam) (*model.Channel, string, *types.NewAPIError) { | |||
| if info.ChannelMeta == nil { | |||
| autoBan := c.GetBool("auto_ban") | |||
| autoBanInt := 1 | |||
| @@ -307,24 +338,41 @@ func getChannel(c *gin.Context, info *relaycommon.RelayInfo, retryParam *service | |||
| Type: c.GetInt("channel_type"), | |||
| Name: c.GetString("channel_name"), | |||
| AutoBan: &autoBanInt, | |||
| }, nil | |||
| }, concreteTaskVideoBindingGroup(c, info.TokenGroup), nil | |||
| } | |||
| channel, selectGroup, err := service.CacheGetRandomSatisfiedChannel(retryParam) | |||
| info.PriceData.GroupRatioInfo = helper.HandleGroupRatio(c, info) | |||
| if err != nil { | |||
| return nil, types.NewError(fmt.Errorf("获取分组 %s 下模型 %s 的可用渠道失败(retry): %s", selectGroup, info.OriginModelName, err.Error()), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry()) | |||
| if apiErr, ok := err.(*types.NewAPIError); ok { | |||
| return nil, selectGroup, apiErr | |||
| } | |||
| return nil, selectGroup, types.NewError(fmt.Errorf("获取分组 %s 下模型 %s 的可用渠道失败(retry): %s", selectGroup, info.OriginModelName, err.Error()), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry()) | |||
| } | |||
| if channel == nil { | |||
| return nil, types.NewError(fmt.Errorf("分组 %s 下模型 %s 的可用渠道不存在(retry)", selectGroup, info.OriginModelName), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry()) | |||
| return nil, selectGroup, types.NewError(fmt.Errorf("分组 %s 下模型 %s 的可用渠道不存在(retry)", selectGroup, info.OriginModelName), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry()) | |||
| } | |||
| newAPIError := middleware.SetupContextForSelectedChannel(c, channel, info.OriginModelName) | |||
| if newAPIError != nil { | |||
| return nil, newAPIError | |||
| return nil, selectGroup, newAPIError | |||
| } | |||
| return channel, nil | |||
| return channel, selectGroup, nil | |||
| } | |||
| func requiredTaskChannelTypeForRequest(c *gin.Context) int { | |||
| return 0 | |||
| } | |||
| func allowedTaskChannelTypesForRequest(c *gin.Context) []int { | |||
| if strings.HasPrefix(c.Request.URL.Path, "/api/v3/contents/generations/tasks") { | |||
| return service.VideoAssetChannelTypesForFamily(service.VideoAssetFamilySeedance) | |||
| } | |||
| if isKlingAipingNativePath(c.Request.URL.Path) { | |||
| return service.VideoAssetChannelTypesForFamily(service.VideoAssetFamilyKling) | |||
| } | |||
| return nil | |||
| } | |||
| func shouldRetry(c *gin.Context, openaiErr *types.NewAPIError, retryTimes int) bool { | |||
| @@ -488,6 +536,101 @@ func RelayTaskFetch(c *gin.Context) { | |||
| } | |||
| } | |||
| func preloadTaskPricingConfig(c *gin.Context, info *relaycommon.RelayInfo) *dto.TaskError { | |||
| modelName := strings.TrimSpace(info.OriginModelName) | |||
| action := strings.TrimSpace(info.Action) | |||
| contentType := c.Request.Header.Get("Content-Type") | |||
| storage, err := common.GetBodyStorage(c) | |||
| if err != nil { | |||
| status := http.StatusBadRequest | |||
| if common.IsRequestBodyTooLargeError(err) || errors.Is(err, common.ErrRequestBodyTooLarge) { | |||
| status = http.StatusRequestEntityTooLarge | |||
| } | |||
| return service.TaskErrorWrapperLocal(err, "read_request_body_failed", status) | |||
| } | |||
| defer func() { | |||
| _, _ = storage.Seek(0, io.SeekStart) | |||
| c.Request.Body = io.NopCloser(storage) | |||
| }() | |||
| switch { | |||
| case strings.HasPrefix(contentType, "application/json"): | |||
| body, err := storage.Bytes() | |||
| if err != nil { | |||
| return service.TaskErrorWrapperLocal(err, "read_request_body_failed", http.StatusBadRequest) | |||
| } | |||
| if strings.TrimSpace(string(body)) != "" { | |||
| var payload map[string]any | |||
| if err := common.Unmarshal(body, &payload); err != nil { | |||
| return service.TaskErrorWrapperLocal(err, "invalid_json", http.StatusBadRequest) | |||
| } | |||
| if raw, ok := payload["model_name"].(string); ok && strings.TrimSpace(raw) != "" { | |||
| modelName = strings.TrimSpace(raw) | |||
| } else if raw, ok := payload["model"].(string); ok && strings.TrimSpace(raw) != "" { | |||
| modelName = strings.TrimSpace(raw) | |||
| } | |||
| if raw, ok := payload["action"].(string); ok && strings.TrimSpace(raw) != "" { | |||
| action = strings.TrimSpace(raw) | |||
| } | |||
| } | |||
| case strings.Contains(contentType, gin.MIMEMultipartPOSTForm): | |||
| form, err := common.ParseMultipartFormReusable(c) | |||
| if err != nil { | |||
| return service.TaskErrorWrapperLocal(err, "invalid_multipart_form", http.StatusBadRequest) | |||
| } | |||
| if vals := form.Value["model"]; len(vals) > 0 && strings.TrimSpace(vals[0]) != "" { | |||
| modelName = strings.TrimSpace(vals[0]) | |||
| } | |||
| if vals := form.Value["action"]; len(vals) > 0 && strings.TrimSpace(vals[0]) != "" { | |||
| action = strings.TrimSpace(vals[0]) | |||
| } | |||
| } | |||
| if modelName == "" && action != "" { | |||
| platform := constant.TaskPlatform(c.GetString("platform")) | |||
| modelName = service.CoverTaskActionToModelName(platform, action) | |||
| } | |||
| if modelName != "" { | |||
| info.OriginModelName = modelName | |||
| } | |||
| if action != "" { | |||
| info.Action = action | |||
| } | |||
| info.PricingConfigSnapshotLoaded = true | |||
| info.PricingConfigSnapshot = ratio_setting.GetPricingConfig(info.OriginModelName) | |||
| info.RequireMatrixUsageBilling = info.PricingConfigSnapshot != nil && | |||
| info.PricingConfigSnapshot.BillingUnit == types.BillingUnitPer1MTokens | |||
| return nil | |||
| } | |||
| func buildTaskBillingContext(info *relaycommon.RelayInfo) *model.TaskBillingContext { | |||
| bc := &model.TaskBillingContext{ | |||
| ModelPrice: info.PriceData.ModelPrice, | |||
| GroupRatio: info.PriceData.GroupRatioInfo.GroupRatio, | |||
| ModelRatio: info.PriceData.ModelRatio, | |||
| OtherRatios: info.PriceData.OtherRatios, | |||
| OriginModelName: info.OriginModelName, | |||
| PerCallBilling: common.StringsContains(constant.TaskPricePatches, info.OriginModelName), | |||
| } | |||
| if decision := info.PricingDecisionFrozen; decision != nil && decision.BillingMode == types.BillingModeMatrix { | |||
| bc.BillingMode = decision.BillingMode | |||
| bc.PricingSnapshot = types.CloneMapAny(decision.Snapshot) | |||
| bc.BillingUnit = decision.BillingUnit | |||
| bc.TokenUnitPriceUSD = decision.TokenUnitPriceUSD | |||
| bc.PerCallBilling = decision.PerCallBilling | |||
| if decision.BillingUnit == types.BillingUnitPer1MTokens { | |||
| bc.ModelPrice = decision.TokenUnitPriceUSD | |||
| } else { | |||
| bc.ModelPrice = decision.PriceUSD | |||
| } | |||
| bc.GroupRatio = decision.GroupRatioInfo.GroupRatio | |||
| bc.OtherRatios = types.CloneRatios(decision.OtherRatios) | |||
| } | |||
| return bc | |||
| } | |||
| func RelayTask(c *gin.Context) { | |||
| relayInfo, err := relaycommon.GenRelayInfo(c, types.RelayFormatTask, nil, nil) | |||
| if err != nil { | |||
| @@ -498,11 +641,18 @@ func RelayTask(c *gin.Context) { | |||
| }) | |||
| return | |||
| } | |||
| relayTaskWithInfo(c, relayInfo) | |||
| } | |||
| func relayTaskWithInfo(c *gin.Context, relayInfo *relaycommon.RelayInfo) { | |||
| if taskErr := relay.ResolveOriginTask(c, relayInfo); taskErr != nil { | |||
| respondTaskError(c, taskErr) | |||
| return | |||
| } | |||
| if taskErr := preloadTaskPricingConfig(c, relayInfo); taskErr != nil { | |||
| respondTaskError(c, taskErr) | |||
| return | |||
| } | |||
| var result *relay.TaskSubmitResult | |||
| var taskErr *dto.TaskError | |||
| @@ -513,26 +663,28 @@ func RelayTask(c *gin.Context) { | |||
| }() | |||
| retryParam := &service.RetryParam{ | |||
| Ctx: c, | |||
| TokenGroup: relayInfo.TokenGroup, | |||
| ModelName: relayInfo.OriginModelName, | |||
| Retry: common.GetPointer(0), | |||
| Ctx: c, | |||
| TokenGroup: relayInfo.TokenGroup, | |||
| ModelName: relayInfo.OriginModelName, | |||
| Retry: common.GetPointer(0), | |||
| RequiredChannelType: requiredTaskChannelTypeForRequest(c), | |||
| AllowedChannelTypes: allowedTaskChannelTypesForRequest(c), | |||
| } | |||
| for ; retryParam.GetRetry() <= common.RetryTimes; retryParam.IncreaseRetry() { | |||
| var channel *model.Channel | |||
| var selectedGroup string | |||
| if lockedCh, ok := relayInfo.LockedChannel.(*model.Channel); ok && lockedCh != nil { | |||
| channel = lockedCh | |||
| if retryParam.GetRetry() > 0 { | |||
| if setupErr := middleware.SetupContextForSelectedChannel(c, channel, relayInfo.OriginModelName); setupErr != nil { | |||
| taskErr = service.TaskErrorWrapperLocal(setupErr.Err, "setup_locked_channel_failed", http.StatusInternalServerError) | |||
| break | |||
| } | |||
| selectedGroup = concreteTaskVideoBindingGroup(c, relayInfo.TokenGroup) | |||
| if setupErr := middleware.SetupContextForSelectedChannel(c, channel, relayInfo.OriginModelName); setupErr != nil { | |||
| taskErr = service.TaskErrorWrapperLocal(setupErr.Err, "setup_locked_channel_failed", http.StatusInternalServerError) | |||
| break | |||
| } | |||
| } else { | |||
| var channelErr *types.NewAPIError | |||
| channel, channelErr = getChannel(c, relayInfo, retryParam) | |||
| channel, selectedGroup, channelErr = getChannel(c, relayInfo, retryParam) | |||
| if channelErr != nil { | |||
| logger.LogError(c, channelErr.Error()) | |||
| taskErr = service.TaskErrorWrapperLocal(channelErr.Err, "get_channel_failed", http.StatusInternalServerError) | |||
| @@ -540,6 +692,11 @@ func RelayTask(c *gin.Context) { | |||
| } | |||
| } | |||
| if bindErr := persistTaskVideoBindingIfNeeded(c.GetInt("id"), selectedGroup, channel); bindErr != nil { | |||
| taskErr = service.TaskErrorWrapperLocal(bindErr, "bind_task_video_channel_failed", http.StatusServiceUnavailable) | |||
| break | |||
| } | |||
| addUsedChannel(c, channel.Id) | |||
| bodyStorage, bodyErr := common.GetBodyStorage(c) | |||
| if bodyErr != nil { | |||
| @@ -587,14 +744,8 @@ func RelayTask(c *gin.Context) { | |||
| task.PrivateData.BillingSource = relayInfo.BillingSource | |||
| task.PrivateData.SubscriptionId = relayInfo.SubscriptionId | |||
| task.PrivateData.TokenId = relayInfo.TokenId | |||
| task.PrivateData.BillingContext = &model.TaskBillingContext{ | |||
| ModelPrice: relayInfo.PriceData.ModelPrice, | |||
| GroupRatio: relayInfo.PriceData.GroupRatioInfo.GroupRatio, | |||
| ModelRatio: relayInfo.PriceData.ModelRatio, | |||
| OtherRatios: relayInfo.PriceData.OtherRatios, | |||
| OriginModelName: relayInfo.OriginModelName, | |||
| PerCallBilling: common.StringsContains(constant.TaskPricePatches, relayInfo.OriginModelName), | |||
| } | |||
| task.PrivateData.BillingContext = buildTaskBillingContext(relayInfo) | |||
| task.PrivateData.UpstreamRequest = relay.BuildUpstreamRequestSnapshotForTask(result.UpstreamReqJSON) | |||
| task.Quota = result.Quota | |||
| task.Data = result.TaskData | |||
| task.Action = relayInfo.Action | |||
| @@ -613,9 +764,34 @@ func respondTaskError(c *gin.Context, taskErr *dto.TaskError) { | |||
| if taskErr.StatusCode == http.StatusTooManyRequests { | |||
| taskErr.Message = "当前分组上游负载已饱和,请稍后再试" | |||
| } | |||
| if isKlingAipingNativePath(c.Request.URL.Path) { | |||
| normalizeKlingAipingTaskError(taskErr) | |||
| } | |||
| c.JSON(taskErr.StatusCode, taskErr) | |||
| } | |||
| func isKlingAipingNativePath(path string) bool { | |||
| return pathMatchesAnyKlingAipingNativePrefix(path, | |||
| "/v1/videos/text2video", | |||
| "/v1/videos/image2video", | |||
| "/v1/videos/motion-control", | |||
| "/v1/videos/omni-video", | |||
| "/v1/videos/multi-image2video", | |||
| "/v1/videos/video-extend", | |||
| "/v1/general/advanced-custom-elements", | |||
| "/v1/general/custom-voices", | |||
| ) | |||
| } | |||
| func pathMatchesAnyKlingAipingNativePrefix(path string, prefixes ...string) bool { | |||
| for _, prefix := range prefixes { | |||
| if path == prefix || strings.HasPrefix(path, prefix+"/") { | |||
| return true | |||
| } | |||
| } | |||
| return false | |||
| } | |||
| func shouldRetryTaskRelay(c *gin.Context, channelId int, taskErr *dto.TaskError, retryTimes int) bool { | |||
| if taskErr == nil { | |||
| return false | |||
| @@ -843,12 +843,19 @@ func CreateUser(c *gin.Context) { | |||
| Password: user.Password, | |||
| DisplayName: user.DisplayName, | |||
| Role: user.Role, // 保持管理员设置的角色 | |||
| Email: user.Email, | |||
| Quota: user.Quota, | |||
| Group: user.Group, | |||
| AffCode: user.AffCode, | |||
| } | |||
| if err := cleanUser.Insert(0); err != nil { | |||
| common.ApiError(c, err) | |||
| return | |||
| } | |||
| // 同步用户到海外节点(异步执行,不阻塞创建流程) | |||
| region_sync.PushUserCreateToSlave(&cleanUser) | |||
| c.JSON(http.StatusOK, gin.H{ | |||
| "success": true, | |||
| "message": "", | |||
| @@ -0,0 +1,89 @@ | |||
| # CN 线上 Seedance 任务 — Aiping 上游返回数据 | |||
| > 数据来源:CN 服务器 `root@123.57.74.135`,MySQL 容器 `ov-mysql`,数据库 `new-api`,表 `tasks` | |||
| > 查询时间:2026-07-06 | |||
| ## 任务完整参数 | |||
| | # | task_id | 模型 | aiping_id | 上游 task_id | Total Tokens | 分辨率 | 时长 | 比例 | FPS | Seed | Audio | Tier | | |||
| |---|---------|------|-----------|-------------|:-----------:|:-----:|:---:|:----:|:---:|:----:|:-----:|:----:| | |||
| | 1 | `task_iIyaivFD1nn0h3pO9BvWVb5Z3csNzcj4` | Seedance-2.0 | `ee43bb3f-af2f-4c22-a7e2-cfbf6ff2da01` | `cgt-20260623152942-bvb94` | 108,900 | 720p | 5s | 16:9 | 24 | 23651 | ✅ | default | | |||
| | 2 | `task_8jNrKcgTSKZPjrs2uLhnYNk7WhHxnQW8` | Seedance-2.0 | `daf418a3-32b6-4083-afe9-4dd958396c59` | `cgt-20260623152943-vf5vd` | **245,025** | **1080p** | 5s | 16:9 | 24 | 36954 | ✅ | default | | |||
| | 3 | `task_2a0nxUWALYZuRndTcAQ9XS2alTQ2AeFX` | Seedance-2.0 | `ee91adc3-445f-4359-b5c7-7cd6b482e88d` | `cgt-20260623152943-f4nqx` | **50,638** | **480p** | 5s | 16:9 | 24 | 8813 | ✅ | default | | |||
| | 4 | `task_b1LIHZ6qLjkzZYsxFLcTbVA1rPWKgqDe` | Seedance-2.0-fast | `19cd2e2f-b1fd-4b6d-b2b8-cc21757bb61c` | `cgt-20260623152944-jtbq5` | 108,900 | 720p | 5s | **9:16** | 24 | 6689 | ✅ | default | | |||
| | 5 | `task_lonxGyh9CCLzbr6PU4XTdC3mYjq3aErx` | Seedance-2.0 | `947033ed-0b6d-4eb3-a48e-f80648e98c37` | `cgt-20260623153543-glrq7` | 108,900 | 720p | 5s | 16:9 | 24 | 65323 | ✅ | default | | |||
| | 6 | `task_6IpuBtvZuaLBsHwQcR1tGCncw5rHu2Et` | Seedance-2.0 | `1daf3f8d-d718-4990-bc19-f917457cfdeb` | `cgt-20260623153741-4h5r5` | 108,900 | 720p | 5s | 16:9 | 24 | 4233 | ✅ | default | | |||
| | 7 | `task_paCqHNnxqjl5fLlHjptD8cjS1TXNrE7K` | Seedance-2.0 | `b7ee658f-acac-433d-897d-5c9ed55ef402` | `cgt-20260624142609-vnwrl` | 108,900 | 720p | 5s | 16:9 | 24 | 78491 | ❌ | default | | |||
| | 8 | `task_UGMvicub7MIJSGHIYJ38o2T2Z8ULPyIJ` | Seedance-2.0 | `3828ab25-4ca1-4dd5-bc22-58bbfd723ae4` | `cgt-20260624142609-wt27b` | 108,900 | 720p | 5s | 16:9 | 24 | 62735 | ✅ | default | | |||
| | 9 | `task_jLgr8WpFOS1R8vxXPgYPQAZjMYypG683` | Seedance-2.0-fast | `fda11914-bbce-4937-b601-5a5ac082c05d` | `cgt-20260624142610-j98vj` | 108,900 | 720p | 5s | 16:9 | 24 | 25865 | ❌ | default | | |||
| | 10 | `task_XvJL39JVdEFT28DMcNyAUWoKzOGOtohj` | Seedance-2.0-fast | `0a476ef1-d4de-4364-8050-03261a890c23` | `cgt-20260624142611-kswz9` | 108,900 | 720p | 5s | 16:9 | 24 | 93047 | ✅ | default | | |||
| | 11 | `task_LRvXoCyhbjytg7nVqxi8SHVtjy14ICdU` | Seedance-2.0 | `5e1a7048-d88a-44ea-bba8-a6e0466eaa4a` | `cgt-20260624151204-l95pj` | 108,900 | 720p | 5s | 16:9 | 24 | 97198 | ❌ | default | | |||
| | 12 | `task_dDYpwuJr5GsjfHv1bCPSU2ep2QP6Oxwv` | Seedance-2.0-fast | `58ea6453-30d8-484d-a1ec-f223012a1cb2` | `cgt-20260625112024-dd74d` | 108,900 | 720p | 5s | 16:9 | 24 | 48027 | ❌ | default | | |||
| ## 统计分布 | |||
| ### 模型分布 | |||
| | 模型 | 数量 | 占比 | | |||
| |------|:---:|:---:| | |||
| | Doubao-Seedance-2.0 | 8 | 66.7% | | |||
| | Doubao-Seedance-2.0-fast | 4 | 33.3% | | |||
| ### 分辨率分布 | |||
| | 分辨率 | 数量 | | |||
| |--------|:---:| | |||
| | 480p | 1 | | |||
| | 720p | 10 | | |||
| | 1080p | 1 | | |||
| ### Token 与分辨率关系 | |||
| | 分辨率 | Total Tokens | 倍数 (vs 720p) | | |||
| |--------|:-----------:|:--------------:| | |||
| | 480p | 50,638 | ~0.46× | | |||
| | 720p | 108,900 | 1× (基准) | | |||
| | 1080p | 245,025 | ~2.25× | | |||
| ### 其他参数 | |||
| | 参数 | 分布 | | |||
| |------|------| | |||
| | 时长 | 全部 5s | | |||
| | 比例 | 16:9: 11 / 9:16: 1(竖屏) | | |||
| | FPS | 全部 24 | | |||
| | Audio | ✅ 生成: 8 / ❌ 不生成: 4 | | |||
| | Tier | 全部 default | | |||
| ## 上游响应示例(task #1) | |||
| ```json | |||
| { | |||
| "aiping_id": "ee43bb3f-af2f-4c22-a7e2-cfbf6ff2da01", | |||
| "id": "cgt-20260623152942-bvb94", | |||
| "model": "doubao-seedance-2-0-260128", | |||
| "status": "succeeded", | |||
| "content": { | |||
| "video_url": "https://ark-acg-cn-beijing.tos-cn-beijing.volces.com/..." | |||
| }, | |||
| "duration": 5, | |||
| "resolution": "720p", | |||
| "ratio": "16:9", | |||
| "seed": 23651, | |||
| "framespersecond": 24, | |||
| "service_tier": "default", | |||
| "generate_audio": true, | |||
| "execution_expires_after": 172800, | |||
| "usage": { | |||
| "total_tokens": 108900, | |||
| "completion_tokens": 108900 | |||
| } | |||
| } | |||
| ``` | |||
| ## 关键发现 | |||
| 1. **aiping_id**:上游返回的 UUID 格式标识符,被 `doubao_aiping/adaptor.go:178` 在返给客户端前删除,仅保存在 `tasks.data` 中 | |||
| 2. **上游 task_id**:`cgt-{YYYYMMDDHHmmss}-{5位随机}` 格式 | |||
| 3. **计费维度**:`total_tokens` 严格按分辨率变化 — 720p=108,900 是基准,480p 约减半,1080p 约翻 2.25 倍 | |||
| 4. **模型名映射**:上游返回原始名 `doubao-seedance-2-0-260128`,properties 中存储展示名 `Doubao-Seedance-2.0` | |||
| @@ -0,0 +1,608 @@ | |||
| # E2E Test Foundation Implementation Plan | |||
| > **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. | |||
| **Goal:** Build a reusable Playwright E2E foundation and cover the model multidimensional pricing save/reload workflow. | |||
| **Architecture:** Playwright runs from `web/`, starts the React app through Vite, and starts the Go backend through a global setup helper. The backend uses a temporary SQLite database and a fixed test session secret so each E2E run is isolated and repeatable. | |||
| **Tech Stack:** Playwright, Bun, React/Vite, Go backend, SQLite. | |||
| --- | |||
| ## File Structure | |||
| - Create `web/playwright.config.js`: Playwright configuration, projects, web server, artifacts, and env wiring. | |||
| - Create `web/e2e/global-setup.js`: start backend with temporary SQLite and wait for readiness. | |||
| - Create `web/e2e/global-teardown.js`: stop backend process and remove temporary database directory. | |||
| - Create `web/e2e/auth.setup.js`: log in as `root / 123456` and save storage state. | |||
| - Create `web/e2e/model-pricing.spec.js`: first real E2E workflow for model multidimensional pricing. | |||
| - Create `web/e2e/utils/process.js`: cross-platform process spawn/kill helpers. | |||
| - Create `web/e2e/utils/wait.js`: HTTP polling helper. | |||
| - Modify `web/package.json`: add E2E scripts. | |||
| - Modify `web/src/components/settings/RatioSetting.jsx`: add `data-testid` to the model pricing tab content. | |||
| - Modify `web/src/pages/Setting/Operation/ModelPricing/index.jsx`: add stable selectors for the page, add button, search input, and table. | |||
| - Modify `web/src/pages/Setting/Operation/ModelPricing/Editor.jsx`: add stable selectors for model input, save, generator, and JSON editor. | |||
| - Modify `web/src/pages/Setting/Operation/ModelPricing/DimensionEditor.jsx`: add stable selectors for add dimension and editable cells. | |||
| - Modify `web/src/pages/Setting/Operation/ModelPricing/GeneratorModal.jsx`: add stable selectors for base price, add ratio, apply, and ratio inputs. | |||
| ## Task 1: Add Playwright Scripts | |||
| **Files:** | |||
| - Modify: `web/package.json` | |||
| - [ ] **Step 1: Add E2E scripts** | |||
| Add these entries inside `scripts`: | |||
| ```json | |||
| "e2e": "playwright test", | |||
| "e2e:ui": "playwright test --ui", | |||
| "e2e:headed": "playwright test --headed", | |||
| "e2e:install": "playwright install chromium" | |||
| ``` | |||
| - [ ] **Step 2: Verify package scripts parse** | |||
| Run: | |||
| ```powershell | |||
| cd web | |||
| bun pm pkg get scripts.e2e | |||
| ``` | |||
| Expected output includes: | |||
| ```text | |||
| "playwright test" | |||
| ``` | |||
| - [ ] **Step 3: Commit** | |||
| ```powershell | |||
| git add web/package.json | |||
| git commit -m "test: add e2e scripts" | |||
| ``` | |||
| ## Task 2: Add Backend Process Helpers | |||
| **Files:** | |||
| - Create: `web/e2e/utils/process.js` | |||
| - Create: `web/e2e/utils/wait.js` | |||
| - [ ] **Step 1: Create process helper** | |||
| Create `web/e2e/utils/process.js`: | |||
| ```js | |||
| import { spawn } from 'node:child_process'; | |||
| export function startProcess(command, args, options = {}) { | |||
| const child = spawn(command, args, { | |||
| cwd: options.cwd, | |||
| env: { ...process.env, ...(options.env || {}) }, | |||
| shell: process.platform === 'win32', | |||
| stdio: ['ignore', 'pipe', 'pipe'], | |||
| }); | |||
| const output = []; | |||
| child.stdout.on('data', (chunk) => output.push(chunk.toString())); | |||
| child.stderr.on('data', (chunk) => output.push(chunk.toString())); | |||
| child.once('exit', (code) => { | |||
| if (code !== null && code !== 0) { | |||
| output.push(`\nprocess exited with code ${code}\n`); | |||
| } | |||
| }); | |||
| return { child, output }; | |||
| } | |||
| export async function stopProcess(child) { | |||
| if (!child || child.killed) return; | |||
| await new Promise((resolve) => { | |||
| child.once('exit', resolve); | |||
| child.kill(process.platform === 'win32' ? undefined : 'SIGTERM'); | |||
| setTimeout(() => { | |||
| if (!child.killed) child.kill('SIGKILL'); | |||
| resolve(); | |||
| }, 5000).unref(); | |||
| }); | |||
| } | |||
| ``` | |||
| - [ ] **Step 2: Create HTTP wait helper** | |||
| Create `web/e2e/utils/wait.js`: | |||
| ```js | |||
| export async function waitForHttp(url, options = {}) { | |||
| const timeoutMs = options.timeoutMs || 60000; | |||
| const intervalMs = options.intervalMs || 500; | |||
| const start = Date.now(); | |||
| let lastError; | |||
| while (Date.now() - start < timeoutMs) { | |||
| try { | |||
| const response = await fetch(url, { cache: 'no-store' }); | |||
| if (response.ok) return response; | |||
| lastError = new Error(`HTTP ${response.status}`); | |||
| } catch (error) { | |||
| lastError = error; | |||
| } | |||
| await new Promise((resolve) => setTimeout(resolve, intervalMs)); | |||
| } | |||
| throw new Error(`Timed out waiting for ${url}: ${lastError?.message || 'no response'}`); | |||
| } | |||
| ``` | |||
| - [ ] **Step 3: Commit** | |||
| ```powershell | |||
| git add web/e2e/utils/process.js web/e2e/utils/wait.js | |||
| git commit -m "test: add e2e process helpers" | |||
| ``` | |||
| ## Task 3: Add Playwright Runtime Setup | |||
| **Files:** | |||
| - Create: `web/e2e/global-setup.js` | |||
| - Create: `web/e2e/global-teardown.js` | |||
| - Create: `web/playwright.config.js` | |||
| - [ ] **Step 1: Create global setup** | |||
| Create `web/e2e/global-setup.js`: | |||
| ```js | |||
| import fs from 'node:fs/promises'; | |||
| import os from 'node:os'; | |||
| import path from 'node:path'; | |||
| import { startProcess } from './utils/process.js'; | |||
| import { waitForHttp } from './utils/wait.js'; | |||
| const statePath = path.resolve('.playwright/e2e-state.json'); | |||
| export default async function globalSetup() { | |||
| const repoRoot = path.resolve('..'); | |||
| const backendPort = process.env.E2E_BACKEND_PORT || '3001'; | |||
| const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), 'new-api-e2e-')); | |||
| const sqlitePath = path.join(tempDir, 'e2e.db'); | |||
| await fs.mkdir(path.dirname(statePath), { recursive: true }); | |||
| const backend = startProcess('go', ['run', 'main.go'], { | |||
| cwd: repoRoot, | |||
| env: { | |||
| PORT: backendPort, | |||
| SQLITE_PATH: `${sqlitePath}?_busy_timeout=30000`, | |||
| SESSION_SECRET: 'e2e_session_secret_for_local_tests', | |||
| }, | |||
| }); | |||
| try { | |||
| await waitForHttp(`http://127.0.0.1:${backendPort}/api/setup`, { timeoutMs: 90000 }); | |||
| } catch (error) { | |||
| console.error(backend.output.join('')); | |||
| backend.child.kill(); | |||
| throw error; | |||
| } | |||
| await fs.writeFile( | |||
| statePath, | |||
| JSON.stringify({ | |||
| backendPid: backend.child.pid, | |||
| tempDir, | |||
| backendPort, | |||
| }), | |||
| 'utf8', | |||
| ); | |||
| } | |||
| ``` | |||
| - [ ] **Step 2: Create global teardown** | |||
| Create `web/e2e/global-teardown.js`: | |||
| ```js | |||
| import fs from 'node:fs/promises'; | |||
| import path from 'node:path'; | |||
| const statePath = path.resolve('.playwright/e2e-state.json'); | |||
| export default async function globalTeardown() { | |||
| try { | |||
| const state = JSON.parse(await fs.readFile(statePath, 'utf8')); | |||
| if (state.backendPid) { | |||
| try { | |||
| process.kill(state.backendPid); | |||
| } catch {} | |||
| } | |||
| if (state.tempDir) { | |||
| await fs.rm(state.tempDir, { recursive: true, force: true }); | |||
| } | |||
| await fs.rm(statePath, { force: true }); | |||
| } catch {} | |||
| } | |||
| ``` | |||
| - [ ] **Step 3: Create Playwright config** | |||
| Create `web/playwright.config.js`: | |||
| ```js | |||
| import { defineConfig, devices } from '@playwright/test'; | |||
| const frontendPort = process.env.E2E_FRONTEND_PORT || '5173'; | |||
| const backendPort = process.env.E2E_BACKEND_PORT || '3001'; | |||
| export default defineConfig({ | |||
| testDir: './e2e', | |||
| globalSetup: './e2e/global-setup.js', | |||
| globalTeardown: './e2e/global-teardown.js', | |||
| timeout: 60_000, | |||
| expect: { timeout: 10_000 }, | |||
| fullyParallel: false, | |||
| retries: process.env.CI ? 1 : 0, | |||
| reporter: [['list'], ['html', { open: 'never' }]], | |||
| use: { | |||
| baseURL: `http://127.0.0.1:${frontendPort}`, | |||
| trace: 'retain-on-failure', | |||
| screenshot: 'only-on-failure', | |||
| video: 'off', | |||
| }, | |||
| webServer: { | |||
| command: `bun run dev -- --host 127.0.0.1 --port ${frontendPort}`, | |||
| url: `http://127.0.0.1:${frontendPort}`, | |||
| reuseExistingServer: !process.env.CI, | |||
| env: { | |||
| VITE_REACT_APP_SERVER_URL: `http://127.0.0.1:${backendPort}`, | |||
| }, | |||
| timeout: 90_000, | |||
| }, | |||
| projects: [ | |||
| { name: 'setup', testMatch: /auth\.setup\.js/ }, | |||
| { | |||
| name: 'chromium', | |||
| use: { | |||
| ...devices['Desktop Chrome'], | |||
| storageState: 'e2e/.auth/root.json', | |||
| }, | |||
| dependencies: ['setup'], | |||
| testIgnore: /auth\.setup\.js/, | |||
| }, | |||
| ], | |||
| }); | |||
| ``` | |||
| - [ ] **Step 4: Run Playwright list to verify config loads** | |||
| Run: | |||
| ```powershell | |||
| cd web | |||
| bunx playwright test --list | |||
| ``` | |||
| Expected: command exits successfully and lists no feature tests yet, or only setup if Task 4 has already been completed. | |||
| - [ ] **Step 5: Commit** | |||
| ```powershell | |||
| git add web/playwright.config.js web/e2e/global-setup.js web/e2e/global-teardown.js | |||
| git commit -m "test: add e2e runtime setup" | |||
| ``` | |||
| ## Task 4: Add Login Storage Setup | |||
| **Files:** | |||
| - Create: `web/e2e/auth.setup.js` | |||
| - [ ] **Step 1: Create auth setup test** | |||
| Create `web/e2e/auth.setup.js`: | |||
| ```js | |||
| import { test, expect } from '@playwright/test'; | |||
| import fs from 'node:fs/promises'; | |||
| test('authenticate as root', async ({ page }) => { | |||
| await fs.mkdir('e2e/.auth', { recursive: true }); | |||
| await page.goto('/login'); | |||
| await page.getByPlaceholder(/用户名|username/i).fill('root'); | |||
| await page.getByPlaceholder(/密码|password/i).fill('123456'); | |||
| const loginResponse = page.waitForResponse((response) => | |||
| response.url().includes('/api/user/login') && response.request().method() === 'POST', | |||
| ); | |||
| await page.getByRole('button', { name: /登录|login/i }).click(); | |||
| const response = await loginResponse; | |||
| expect(response.ok()).toBeTruthy(); | |||
| await expect(page).toHaveURL(/\/console/); | |||
| await page.context().storageState({ path: 'e2e/.auth/root.json' }); | |||
| }); | |||
| ``` | |||
| - [ ] **Step 2: Run setup only** | |||
| Run: | |||
| ```powershell | |||
| cd web | |||
| bunx playwright test --project=setup | |||
| ``` | |||
| Expected: setup test passes and creates `web/e2e/.auth/root.json`. | |||
| - [ ] **Step 3: Ignore generated auth state** | |||
| If `web/e2e/.auth/root.json` appears in `git status`, add this to `.gitignore`: | |||
| ```gitignore | |||
| web/e2e/.auth/ | |||
| web/.playwright/ | |||
| web/playwright-report/ | |||
| web/test-results/ | |||
| ``` | |||
| - [ ] **Step 4: Commit** | |||
| ```powershell | |||
| git add web/e2e/auth.setup.js .gitignore | |||
| git commit -m "test: add e2e root login setup" | |||
| ``` | |||
| ## Task 5: Add Stable Selectors To Model Pricing UI | |||
| **Files:** | |||
| - Modify: `web/src/components/settings/RatioSetting.jsx` | |||
| - Modify: `web/src/pages/Setting/Operation/ModelPricing/index.jsx` | |||
| - Modify: `web/src/pages/Setting/Operation/ModelPricing/Editor.jsx` | |||
| - Modify: `web/src/pages/Setting/Operation/ModelPricing/DimensionEditor.jsx` | |||
| - Modify: `web/src/pages/Setting/Operation/ModelPricing/GeneratorModal.jsx` | |||
| - [ ] **Step 1: Add tab content selector** | |||
| In `web/src/components/settings/RatioSetting.jsx`, wrap the model pricing tab content: | |||
| ```jsx | |||
| <Tabs.TabPane tab={t('模型多维计费')} itemKey='model_pricing'> | |||
| <div data-testid='model-pricing-tab'> | |||
| <ModelPricing /> | |||
| </div> | |||
| </Tabs.TabPane> | |||
| ``` | |||
| - [ ] **Step 2: Add page-level selectors** | |||
| In `web/src/pages/Setting/Operation/ModelPricing/index.jsx`, use these props: | |||
| ```jsx | |||
| <Card style={{ marginTop: 10 }} data-testid='model-pricing-page'> | |||
| ``` | |||
| ```jsx | |||
| <Button data-testid='model-pricing-add-model' icon={<Plus size={16} />} type='primary' onClick={() => setEditingModel('')}> | |||
| ``` | |||
| ```jsx | |||
| <Input data-testid='model-pricing-search' placeholder={t('搜索模型')} value={keyword} onChange={setKeyword} style={{ width: 280 }} /> | |||
| ``` | |||
| ```jsx | |||
| <Table data-testid='model-pricing-table' columns={columns} dataSource={data} rowKey='model' pagination={{ pageSize: 10 }} loading={loading} /> | |||
| ``` | |||
| - [ ] **Step 3: Add editor selectors** | |||
| In `web/src/pages/Setting/Operation/ModelPricing/Editor.jsx`, add: | |||
| ```jsx | |||
| <SideSheet | |||
| data-testid='model-pricing-editor' | |||
| ``` | |||
| ```jsx | |||
| <Button data-testid='model-pricing-save' type='primary' icon={<Save size={16} />} onClick={save} loading={saving}> | |||
| ``` | |||
| ```jsx | |||
| <Input data-testid='model-pricing-model-name' prefix={t('模型')} value={name} onChange={setName} style={{ width: 260 }} disabled={Boolean(modelName)} /> | |||
| ``` | |||
| ```jsx | |||
| <Button data-testid='model-pricing-open-generator' icon={<Wand2 size={16} />} onClick={() => setGeneratorVisible(true)}> | |||
| ``` | |||
| ```jsx | |||
| <TextArea data-testid='model-pricing-json-editor' value={jsonText} onChange={setJsonText} autosize={{ minRows: 18, maxRows: 28 }} /> | |||
| ``` | |||
| - [ ] **Step 4: Add dimension selectors** | |||
| In `web/src/pages/Setting/Operation/ModelPricing/DimensionEditor.jsx`, add indexed `data-testid` attributes to editable inputs: | |||
| ```jsx | |||
| <Input data-testid={`dimension-key-${index}`} value={value} onChange={(key) => update(index, { key })} /> | |||
| ``` | |||
| ```jsx | |||
| <Input | |||
| data-testid={`dimension-source-path-${index}`} | |||
| value={parsed.path} | |||
| placeholder={parsed.sourceType === DIMENSION_SOURCE_TYPES.advanced ? 'request.options.size' : 'duration'} | |||
| onChange={(path) => update(index, { source: composeDimensionSource(parsed.sourceType, path) })} | |||
| /> | |||
| ``` | |||
| ```jsx | |||
| <Button data-testid='dimension-add' icon={<Plus size={16} />} onClick={add}> | |||
| ``` | |||
| - [ ] **Step 5: Add generator selectors** | |||
| In `web/src/pages/Setting/Operation/ModelPricing/GeneratorModal.jsx`, add: | |||
| ```jsx | |||
| <Modal | |||
| data-testid='model-pricing-generator' | |||
| ``` | |||
| ```jsx | |||
| <InputNumber | |||
| data-testid='generator-base-price' | |||
| ``` | |||
| ```jsx | |||
| <Button data-testid='generator-add-ratio' icon={<Plus size={16} />} onClick={addEntry}> | |||
| ``` | |||
| ```jsx | |||
| <Button data-testid='generator-apply' type='primary' onClick={apply} disabled={preview.rows.length === 0}> | |||
| ``` | |||
| - [ ] **Step 6: Run frontend build** | |||
| Run: | |||
| ```powershell | |||
| cd web | |||
| bun run build | |||
| ``` | |||
| Expected: build exits 0. Existing Vite chunk warnings are acceptable. | |||
| - [ ] **Step 7: Commit** | |||
| ```powershell | |||
| git add web/src/components/settings/RatioSetting.jsx web/src/pages/Setting/Operation/ModelPricing/index.jsx web/src/pages/Setting/Operation/ModelPricing/Editor.jsx web/src/pages/Setting/Operation/ModelPricing/DimensionEditor.jsx web/src/pages/Setting/Operation/ModelPricing/GeneratorModal.jsx | |||
| git commit -m "test: add model pricing e2e selectors" | |||
| ``` | |||
| ## Task 6: Add Model Pricing E2E Workflow | |||
| **Files:** | |||
| - Create: `web/e2e/model-pricing.spec.js` | |||
| - [ ] **Step 1: Write the E2E test** | |||
| Create `web/e2e/model-pricing.spec.js`: | |||
| ```js | |||
| import { test, expect } from '@playwright/test'; | |||
| test('model multidimensional pricing can be saved and loaded', async ({ page, request }) => { | |||
| const modelName = `e2e-video-pricing-model-${Date.now()}`; | |||
| await page.goto('/console/setting?tab=ratio'); | |||
| await page.getByTestId('model-pricing-tab').waitFor(); | |||
| await page.getByTestId('model-pricing-add-model').click(); | |||
| await page.getByTestId('model-pricing-editor').waitFor(); | |||
| await page.getByTestId('model-pricing-model-name').fill(modelName); | |||
| await page.getByTestId('dimension-key-0').fill('duration'); | |||
| await page.getByTestId('dimension-source-path-0').fill('duration'); | |||
| await page.getByTestId('dimension-add').click(); | |||
| await page.getByTestId('dimension-key-1').fill('quality'); | |||
| await page.getByTestId('dimension-source-path-1').fill('quality'); | |||
| await page.getByTestId('model-pricing-open-generator').click(); | |||
| await page.getByTestId('model-pricing-generator').waitFor(); | |||
| await page.getByTestId('generator-base-price').fill('0.2'); | |||
| await page.getByTestId('generator-apply').click(); | |||
| const saveResponsePromise = page.waitForResponse((response) => | |||
| response.url().includes(`/api/option/model_pricing/${encodeURIComponent(modelName)}`) && | |||
| response.request().method() === 'PUT', | |||
| ); | |||
| await page.getByTestId('model-pricing-save').click(); | |||
| const saveResponse = await saveResponsePromise; | |||
| expect(saveResponse.ok()).toBeTruthy(); | |||
| const saveBody = await saveResponse.json(); | |||
| expect(saveBody.success).toBe(true); | |||
| await page.reload(); | |||
| await page.getByTestId('model-pricing-search').fill(modelName); | |||
| await expect(page.getByTestId('model-pricing-table')).toContainText(modelName); | |||
| const apiResponse = await request.get(`/api/option/model_pricing/${encodeURIComponent(modelName)}`); | |||
| expect(apiResponse.ok()).toBeTruthy(); | |||
| const body = await apiResponse.json(); | |||
| expect(body.success).toBe(true); | |||
| expect(body.data.schema_version).toBe(1); | |||
| expect(body.data.scope).toBe('model'); | |||
| expect(body.data.dimensions.map((item) => item.key)).toEqual(['duration', 'quality']); | |||
| expect(body.data.table.length).toBeGreaterThan(0); | |||
| }); | |||
| ``` | |||
| - [ ] **Step 2: Run E2E** | |||
| Run: | |||
| ```powershell | |||
| cd web | |||
| bun run e2e | |||
| ``` | |||
| Expected: setup project logs in, chromium project runs the model pricing workflow, and all tests pass. | |||
| - [ ] **Step 3: Commit** | |||
| ```powershell | |||
| git add web/e2e/model-pricing.spec.js | |||
| git commit -m "test: cover model pricing e2e workflow" | |||
| ``` | |||
| ## Task 7: Final Verification | |||
| **Files:** | |||
| - No new files. | |||
| - [ ] **Step 1: Run model pricing unit tests** | |||
| Run: | |||
| ```powershell | |||
| cd web | |||
| bun test src/pages/Setting/Operation/ModelPricing/pricingConfig.test.js | |||
| ``` | |||
| Expected: all tests pass. | |||
| - [ ] **Step 2: Run frontend build** | |||
| Run: | |||
| ```powershell | |||
| cd web | |||
| bun run build | |||
| ``` | |||
| Expected: build exits 0. Existing chunk warnings are acceptable. | |||
| - [ ] **Step 3: Run E2E** | |||
| Run: | |||
| ```powershell | |||
| cd web | |||
| bun run e2e | |||
| ``` | |||
| Expected: all Playwright projects pass. | |||
| - [ ] **Step 4: Review status** | |||
| Run: | |||
| ```powershell | |||
| git status --short --branch | |||
| git log --oneline -5 | |||
| ``` | |||
| Expected: only intentional changes remain uncommitted, or all E2E commits are present if each task committed successfully. | |||
| @@ -0,0 +1,166 @@ | |||
| # E2E Test Foundation Design | |||
| ## Goal | |||
| Build a reusable end-to-end test foundation for the admin console, then use the model multidimensional pricing page as the first covered workflow. | |||
| The foundation must verify the real frontend/backend contract, authentication, persistence, and page routing. It should be easy to run locally and suitable for CI later. | |||
| ## Scope | |||
| In scope: | |||
| - Add a Playwright-based E2E test foundation under `web/`. | |||
| - Run the React frontend through Vite. | |||
| - Run the Go backend against an isolated SQLite database. | |||
| - Log in as the generated root user. | |||
| - Cover the model multidimensional pricing workflow under `/console/setting?tab=ratio`. | |||
| - Keep tests deterministic by using temporary files, fixed ports, and isolated storage state. | |||
| Out of scope for the first pass: | |||
| - Docker Compose E2E. | |||
| - MySQL/PostgreSQL E2E matrix. | |||
| - Full relay calls to upstream AI providers. | |||
| - Broad admin console coverage outside the model pricing workflow. | |||
| ## Recommended Approach | |||
| Use real backend + temporary SQLite + Playwright. | |||
| The project already has `@playwright/test` in `web/package.json`. The backend can start with `go run main.go`, and when no users exist it creates the root user with username `root` and password `123456`. Running against a temporary SQLite database gives each test run a clean system without requiring external services. | |||
| This approach is slower than API mocks, but it catches the issues that matter for this feature: frontend/backend response shape, save/load behavior, auth redirects, and persistence across page reloads. | |||
| ## Runtime Architecture | |||
| Playwright will run from the `web/` directory. | |||
| Backend: | |||
| - Started by a Playwright global setup helper. | |||
| - Command: `go run main.go` from the repository root. | |||
| - Environment: | |||
| - `PORT=3001` | |||
| - `SQLITE_PATH=<temp-dir>/e2e.db?_busy_timeout=30000` | |||
| - `SESSION_SECRET=e2e_session_secret_for_local_tests` | |||
| - The helper waits for `/api/setup` or another lightweight HTTP endpoint before running tests. | |||
| - The helper stores backend process metadata so global teardown can stop it. | |||
| Frontend: | |||
| - Started by Playwright `webServer`. | |||
| - Command: `bun run dev -- --host 127.0.0.1 --port 5173` | |||
| - Environment: | |||
| - `VITE_REACT_APP_SERVER_URL=http://127.0.0.1:3001` | |||
| - Base URL: `http://127.0.0.1:5173` | |||
| Browser auth: | |||
| - A setup project logs in once as `root / 123456`. | |||
| - The authenticated browser state is saved to `web/e2e/.auth/root.json`. | |||
| - Feature tests reuse that storage state. | |||
| ## Test Layout | |||
| Proposed files: | |||
| - `web/playwright.config.js` | |||
| - `web/e2e/global-setup.js` | |||
| - `web/e2e/global-teardown.js` | |||
| - `web/e2e/auth.setup.js` | |||
| - `web/e2e/model-pricing.spec.js` | |||
| - `web/e2e/utils/process.js` | |||
| - `web/e2e/utils/wait.js` | |||
| Proposed package scripts: | |||
| - `e2e`: `playwright test` | |||
| - `e2e:ui`: `playwright test --ui` | |||
| - `e2e:headed`: `playwright test --headed` | |||
| - `e2e:install`: `playwright install chromium` | |||
| ## First Workflow: Model Multidimensional Pricing | |||
| The first E2E test should cover the critical user path: | |||
| 1. Open `/console/setting?tab=ratio`. | |||
| 2. Confirm the admin setting page loads. | |||
| 3. Open the `模型多维计费` tab. | |||
| 4. Create or edit pricing for a deterministic test model name, for example `e2e-video-pricing-model`. | |||
| 5. Add dimensions such as `duration` and `quality`. | |||
| 6. Use the generator to produce a linear pricing table. | |||
| 7. Save the config. | |||
| 8. Reload the page. | |||
| 9. Reopen the same model config. | |||
| 10. Assert that dimensions and generated rows are loaded. | |||
| The test should also observe the network response for save/load requests and assert successful API envelopes where practical. This makes contract mismatches visible even if the UI shows a generic toast. | |||
| ## Selectors | |||
| Prefer stable selectors instead of brittle text-only selectors for controls that are hard to disambiguate. | |||
| The implementation should add `data-testid` only where current markup cannot be selected reliably: | |||
| - model pricing tab container | |||
| - create/edit model pricing button | |||
| - model name input | |||
| - dimension editor add row button | |||
| - generator open/apply buttons | |||
| - save button | |||
| Visible text can still be used for high-level navigation when the text is unique and user-facing. | |||
| ## Data Isolation | |||
| Each E2E run uses a new temporary SQLite database path. This prevents tests from depending on local developer data and avoids cleanup requirements for most tables. | |||
| Within the UI, test-created models should use an `e2e-` prefix. If the same database is reused during debugging, tests can safely identify their own records. | |||
| ## Failure Artifacts | |||
| Playwright should collect useful artifacts on failure: | |||
| - screenshot: only on failure | |||
| - trace: retain on failure | |||
| - video: off by default for speed | |||
| These defaults keep normal runs fast while still making UI failures diagnosable. | |||
| ## CI Path | |||
| The initial implementation should work locally first. CI can later run: | |||
| 1. `cd web && bun install` | |||
| 2. `cd web && bun run e2e:install` | |||
| 3. `cd web && bun run e2e` | |||
| No external database or service should be required for the first workflow. | |||
| ## Risks And Mitigations | |||
| Backend startup can be slow. | |||
| - Mitigation: wait on an HTTP endpoint with a clear timeout and useful logs. | |||
| Ports can be occupied locally. | |||
| - Mitigation: use fixed defaults first for simplicity, but allow overrides through `E2E_BACKEND_PORT` and `E2E_FRONTEND_PORT`. | |||
| Chinese UI text may be rendered through i18n and encoding-sensitive files. | |||
| - Mitigation: prefer `data-testid` for critical controls and keep user-facing text assertions limited to stable labels. | |||
| Playwright browsers may not be installed. | |||
| - Mitigation: provide `bun run e2e:install` and document it in script names. | |||
| ## Acceptance Criteria | |||
| - `cd web && bun run e2e` starts the backend and frontend automatically. | |||
| - The run uses an isolated SQLite database. | |||
| - The test logs in as root without manual steps. | |||
| - The model multidimensional pricing workflow can save and reload a generated config. | |||
| - Failing tests produce a trace and screenshot. | |||
| - Existing unit tests and frontend build remain runnable without depending on E2E setup. | |||
| @@ -0,0 +1,133 @@ | |||
| # Relay Capture JSON Storage Design | |||
| ## Context | |||
| `RelayCaptureMiddleware` currently writes capture files under `data/relay-capture/YYYY-MM-DD/<request_id>.log` using custom text sections: | |||
| - `=== REQUEST ... ===` | |||
| - request method, path, user ID, headers, and body | |||
| - `=== RESPONSE ===` | |||
| - raw response chunks | |||
| - `=== END duration_ms=... ===` | |||
| This format is easy to append to, but hard to parse reliably. It also stores sensitive request headers in plaintext. The new design changes relay capture output to one structured JSON file per captured request while preserving relay behavior. | |||
| ## Goals | |||
| - Store request headers, request body, response headers, and response body as JSON fields. | |||
| - Generate `data/relay-capture/YYYY-MM-DD/<request_id>.json`. | |||
| - Keep request and response bodies as strings, without JSON parsing. | |||
| - Store response body twice: | |||
| - `response.body` as the complete concatenated response string. | |||
| - `response.body_chunks` as the ordered chunks captured from `ResponseWriter.Write` and `WriteString`. | |||
| - Store headers as `map[string][]string`, matching `http.Header` semantics. | |||
| - Redact sensitive request headers in capture output. | |||
| - Ensure capture failures never change the client-visible relay response. | |||
| - Use `common.Marshal` for JSON encoding. | |||
| ## Non-Goals | |||
| - Do not keep writing the legacy `.log` format. | |||
| - Do not parse request or response bodies into JSON objects. | |||
| - Do not redact request or response bodies in this change. | |||
| - Do not capture upstream provider request or response payloads after adapter conversion. This middleware captures the client-facing relay request and response at the Gin layer. | |||
| - Do not introduce database storage for capture records. | |||
| ## JSON Shape | |||
| Each capture file contains one complete JSON object: | |||
| ```json | |||
| { | |||
| "request_id": "test-non-stream-120000", | |||
| "captured_at": "2026-06-23T12:00:00.000000000Z", | |||
| "duration_ms": 123, | |||
| "user_id": 42, | |||
| "request": { | |||
| "method": "POST", | |||
| "path": "/v1/chat/completions", | |||
| "query": "stream=true", | |||
| "headers": { | |||
| "Authorization": ["[REDACTED]"], | |||
| "Content-Type": ["application/json"] | |||
| }, | |||
| "body": "{\"model\":\"gpt-4\"}" | |||
| }, | |||
| "response": { | |||
| "status_code": 200, | |||
| "headers": { | |||
| "Content-Type": ["application/json; charset=utf-8"] | |||
| }, | |||
| "body": "{\"id\":\"chatcmpl-test\"}", | |||
| "body_chunks": ["{\"id\":\"chatcmpl-test\"}"] | |||
| }, | |||
| "capture_errors": [] | |||
| } | |||
| ``` | |||
| `capture_errors` is present to preserve non-fatal capture problems without failing the relay request. It should be an empty array when capture succeeds. | |||
| ## Middleware Flow | |||
| 1. Check `capture_relay` cache with the authenticated user ID. Disabled users continue through the middleware unchanged. | |||
| 2. Record `start`, `captured_at`, `request_id`, user ID, method, path, query, and a redacted copy of request headers. | |||
| 3. Read the request body through `common.GetBodyStorage(c)`, not directly from `c.Request.Body`, so downstream handlers can continue using the reusable body storage. | |||
| 4. Replace `c.Writer` with a capture response writer that delegates writes to the original writer and copies only the successfully written bytes into an in-memory collector. | |||
| 5. Run `c.Next()`. | |||
| 6. Build response metadata from the final status code, final response headers, collected `body`, and ordered `body_chunks`. | |||
| 7. Marshal the full record with `common.Marshal`. | |||
| 8. Write the JSON file under `data/relay-capture/YYYY-MM-DD/<request_id>.json`. | |||
| If `request_id` is empty, generate a fallback value like `capture-<unix_nano>` to avoid invalid or colliding file names in tests and edge paths. | |||
| ## Header Redaction | |||
| Request header redaction is case-insensitive and only affects the capture record. It must not mutate `c.Request.Header`. | |||
| Initial sensitive request headers: | |||
| - `Authorization` | |||
| - `Proxy-Authorization` | |||
| - `X-Api-Key` | |||
| - `Api-Key` | |||
| - `Cookie` | |||
| Each value for a sensitive header is replaced with `[REDACTED]`. Response headers are copied as-is in this change. | |||
| ## Error Handling | |||
| Capture is diagnostic. It must not block or alter relay behavior. | |||
| - If request body capture fails, continue the request and append a message to `capture_errors`. | |||
| - If response writes fail, return the underlying `ResponseWriter` result exactly as today. The collector only records bytes that were successfully written. | |||
| - If JSON marshaling fails, log with `common.SysError` and do not write a partial capture file. | |||
| - If directory creation or file writing fails, log with `common.SysError` and leave the relay response untouched. | |||
| - The existing writer goroutine and `=== END` text trailer are removed because JSON is written once after `c.Next()`. | |||
| ## Testing | |||
| Update the existing relay capture tests around `middleware/relay_capture.go`. | |||
| Unit coverage: | |||
| - Header copying preserves `map[string][]string` and redacts sensitive request headers, including case variants. | |||
| - Capture file creation uses `.json` and the existing date directory layout. | |||
| - The response capture writer preserves client-visible output while collecting `body` and `body_chunks`. | |||
| - Request body capture stores POST JSON as a string and GET/no-body requests as an empty string. | |||
| Integration coverage: | |||
| - Enabled non-streaming users produce a valid JSON file with request metadata, redacted request headers, response headers, status code, body, `duration_ms`, `user_id`, and `request_id`. | |||
| - Enabled streaming users produce `response.body` containing the full SSE response and `response.body_chunks` preserving chunk order, including `[DONE]`. | |||
| - Disabled users do not produce capture files. | |||
| - Capturing does not change the client-visible response body. | |||
| - Concurrent requests produce one valid JSON file per request. | |||
| - Sensitive header values such as `Bearer sk-test` are absent from capture files and replaced with `[REDACTED]`. | |||
| - Empty request IDs use the fallback `capture-*.json` file name. | |||
| Primary verification command: | |||
| ```powershell | |||
| go test ./middleware | |||
| ``` | |||
| @@ -0,0 +1,322 @@ | |||
| # Model Display Pricing Design | |||
| ## Background | |||
| The current pricing system mixes two responsibilities in one structure: | |||
| - Real billing reads `PricingConfig.table[].price` through relay pricing lookup, freezes it into `PricingDecision`, and later uses that decision for pre-consumption and final settlement. | |||
| - The pricing page also reads `PricingConfig.table[].price` to render model price summaries and matrix details. | |||
| Some video providers need a user-facing estimate such as `768p-6s, price per second`, while the actual charge is settled by token usage, provider-returned deduction, or another backend billing result. Reusing `PricingConfig` for this would make display-only prices affect real billing. | |||
| ## Goal | |||
| Add a display-only pricing configuration that lets the frontend show a supplier-style price table for user estimation. | |||
| This configuration must: | |||
| - Use USD as the base amount so the existing frontend currency conversion can support USD, CNY, and custom currencies. | |||
| - Never participate in quota pre-consumption, final settlement, token billing, task polling, or billing logs. | |||
| - Override only the public pricing page display when configured. | |||
| - Preserve current pricing behavior for models without display pricing. | |||
| ## Non-Goals | |||
| - Do not change how real billing is calculated. | |||
| - Do not add formula-based display pricing in this iteration. | |||
| - Do not replace `ModelPricingRules`; real matrix pricing remains unchanged. | |||
| - Do not calculate exact task cost from display pricing. It is only an estimate. | |||
| ## Data Model | |||
| Add a new option key, separate from `ModelPricingRules`: | |||
| ```text | |||
| ModelDisplayPricing | |||
| ``` | |||
| The option value is a JSON map keyed by model name: | |||
| ```json | |||
| { | |||
| "kling-video": [ | |||
| { | |||
| "specification": "768p-6s", | |||
| "official_supplier_tip": "Text-to-video and image-to-video, 768P 6s", | |||
| "price": 0.33333, | |||
| "unit": "second", | |||
| "sort_order": 1, | |||
| "discount_rate": 1 | |||
| }, | |||
| { | |||
| "specification": "768p-10s", | |||
| "official_supplier_tip": "Text-to-video and image-to-video, 768P 10s", | |||
| "price": 0.4, | |||
| "unit": "second", | |||
| "sort_order": 2, | |||
| "discount_rate": 1 | |||
| }, | |||
| { | |||
| "specification": "1080p", | |||
| "official_supplier_tip": "Text-to-video and image-to-video, 1080P", | |||
| "price": 0.58333, | |||
| "unit": "second", | |||
| "sort_order": 3, | |||
| "discount_rate": 1 | |||
| } | |||
| ] | |||
| } | |||
| ``` | |||
| Suggested Go type: | |||
| ```go | |||
| type ModelDisplayPricingItem struct { | |||
| Specification string `json:"specification"` | |||
| OfficialSupplierTip string `json:"official_supplier_tip"` | |||
| Price float64 `json:"price"` | |||
| Unit string `json:"unit"` | |||
| SortOrder int `json:"sort_order"` | |||
| DiscountRate float64 `json:"discount_rate"` | |||
| } | |||
| ``` | |||
| Field semantics: | |||
| - `specification`: short SKU label shown in compact and detail views. | |||
| - `official_supplier_tip`: supplier-facing explanation shown in detail view. | |||
| - `price`: USD base display price before discount. This is not quota and not billing input. | |||
| - `unit`: display unit, such as `second`, `call`, `image`, or `video`. | |||
| - `sort_order`: stable ordering for the frontend and API response. | |||
| - `discount_rate`: display-only multiplier. `0` or omitted means no discount and is normalized to `1`. | |||
| Validation rules: | |||
| - Model key must not be empty. | |||
| - `specification` must not be empty. | |||
| - `price` must be greater than or equal to `0`. | |||
| - `unit` must not be empty. | |||
| - `discount_rate` is normalized before validation: `0` or omitted becomes `1`; after normalization, the accepted range is `0 < discount_rate <= 1`. | |||
| - Items should be returned sorted by `sort_order`, then `specification`. | |||
| ## Backend Design | |||
| Add a new ratio setting module, parallel to `model_pricing.go`, for display pricing only: | |||
| - `GetModelDisplayPricing(model string) []types.ModelDisplayPricingItem` | |||
| - `GetModelDisplayPricingCopy() map[string][]types.ModelDisplayPricingItem` | |||
| - `SetModelDisplayPricing(model string, items []types.ModelDisplayPricingItem) error` | |||
| - `DeleteModelDisplayPricing(model string) error` | |||
| - `UpdateModelDisplayPricingByJSONString(jsonStr string) error` | |||
| - `ModelDisplayPricing2JSONString() string` | |||
| Store the JSON in `options`, matching the existing config pattern and avoiding a cross-database migration for a display-only setting. | |||
| The option lifecycle must follow the existing `ModelPricingRules` pattern: | |||
| - Add `ModelDisplayPricing` to `model.InitOptionMap()` with an empty JSON object default. | |||
| - Add `case "ModelDisplayPricing"` to `model.updateOptionMap()` so DB-loaded values hydrate the in-memory display pricing cache during startup and option sync. | |||
| - Persist updates through `model.UpdateOption(ModelDisplayPricingOptionKey, ModelDisplayPricing2JSONString())`. | |||
| - Call `model.RefreshPricing()` after successful create, update, or delete so `/api/pricing` reflects the new display data immediately. | |||
| Extend `model.Pricing`: | |||
| ```go | |||
| DisplayPricing []types.ModelDisplayPricingItem `json:"display_pricing,omitempty"` | |||
| ``` | |||
| During `model.updatePricing()`, after the model metadata and real pricing fields are assembled, attach display pricing for the model if configured. This keeps `/api/pricing` as the single source used by the existing pricing page. | |||
| `/api/pricing/user/*model` should also return `display_pricing` when present, but it must not use it for `original_price`, `user_price`, `original_input`, `original_output`, or savings calculation. | |||
| Admin APIs should mirror the existing model pricing rule APIs: | |||
| - `GET /api/option/model_display_pricing` | |||
| - `GET /api/option/model_display_pricing/*model` | |||
| - `PUT /api/option/model_display_pricing/*model` | |||
| - `DELETE /api/option/model_display_pricing/*model` | |||
| Use wildcard model parameters, matching `/api/pricing/user/*model`, so model names containing `/` can be configured without URL-shape bugs. These endpoints manage only display data and call `model.RefreshPricing()` after updates. | |||
| ## Frontend Design | |||
| The display pricing UI is part of the existing model pricing page. It should be implemented as a display override on top of the current `/api/pricing` response, not as a new standalone public page. | |||
| Affected frontend entry points: | |||
| - `web/src/hooks/model-pricing/useModelPricingData.jsx`: keep loading `/api/pricing`; no extra fetch is needed for public display. | |||
| - `web/src/helpers/utils.jsx`: add shared helpers so table, card, and detail views render the same display price. | |||
| - `web/src/components/table/model-pricing/view/table/PricingTableColumns.jsx`: compact price column display. | |||
| - `web/src/components/table/model-pricing/view/card/PricingCardView.jsx`: compact card display. | |||
| - `web/src/components/table/model-pricing/modal/components/ModelPricingTable.jsx`: full detail table display. | |||
| - `web/src/components/table/model-pricing/modal/components/UserPriceComparison.jsx`: hide when display pricing exists. | |||
| Add display pricing helpers in the model pricing utilities: | |||
| - `hasDisplayPricing(record)` | |||
| - `normalizeDisplayPricingItems(record)` | |||
| - `formatDisplayPricingItem({ item, displayPrice, t })` | |||
| - `summarizeDisplayPricing({ record, displayPrice, t })` | |||
| Display calculation: | |||
| ```js | |||
| original = displayPrice(item.price) | |||
| discounted = displayPrice(item.price * item.discount_rate) | |||
| ``` | |||
| If `discount_rate` is `1`, show only the regular price. If lower than `1`, show discounted price and optionally the original price with strikethrough. | |||
| Unit labels should be translated in the frontend. The API returns stable unit keys such as `second`, `call`, `image`, and `video`; the UI maps them to localized labels such as `/ second` or `/ sec` depending on existing i18n conventions. | |||
| Display priority: | |||
| 1. If `record.display_pricing` exists and has items, render display pricing. | |||
| 2. Otherwise, keep the existing matrix pricing, per-token pricing, and per-call pricing behavior. | |||
| ### Pricing Table Display | |||
| In the main table price column, display pricing replaces the current matrix/per-token/per-call price block for that model. | |||
| For multiple display pricing items with the same unit, show a converted price range: | |||
| ```text | |||
| CNY 2.41 - CNY 4.22 / second | |||
| Display price | |||
| ``` | |||
| For one item, or when all items have the same discounted price: | |||
| ```text | |||
| CNY 2.41 / second | |||
| Display price | |||
| ``` | |||
| For multiple units, group by unit and render one short line per unit: | |||
| ```text | |||
| CNY 2.41 - CNY 4.22 / second | |||
| CNY 1.00 / image | |||
| Display price | |||
| ``` | |||
| The compact table view should not render `official_supplier_tip`; that belongs in the detail view. | |||
| ### Pricing Card Display | |||
| Cards should use the same summary output as the table column, with at most the first three display items shown below the range: | |||
| ```text | |||
| CNY 2.41 - CNY 4.22 / second | |||
| 768p-6s CNY 2.41 / second | |||
| 768p-10s CNY 2.90 / second | |||
| 1080p CNY 4.22 / second | |||
| + 2 more | |||
| ``` | |||
| If the card is too narrow on mobile, keep only the range and the `Display price` label. Do not render a horizontal table inside the card. | |||
| ### Detail Side Sheet Display | |||
| When `display_pricing` exists, `ModelPricingTable` should render a dedicated display pricing table instead of the existing group price table or matrix table. | |||
| | Specification | Supplier Tip | Price | Discount | | |||
| | --- | --- | --- | --- | | |||
| | 768p-6s | Text-to-video and image-to-video, 768P 6s | CNY 2.41 / second | - | | |||
| | 768p-10s | Text-to-video and image-to-video, 768P 10s | CNY 2.90 / second | - | | |||
| | 1080p | Text-to-video and image-to-video, 1080P | CNY 4.22 / second | - | | |||
| For discounted items, render the discounted price as the primary value and the original converted price as secondary strikethrough text: | |||
| ```text | |||
| CNY 1.93 / second | |||
| Original CNY 2.41, 20% off | |||
| ``` | |||
| Every display pricing block must include an i18n note with this meaning: | |||
| ```text | |||
| Display price is only an estimate. Actual consumption is based on the final quota deduction after the task completes. | |||
| ``` | |||
| The note must appear in the detail side sheet and in any compact tooltip/popover if one is added later. The main table and card can show only a short `Display price` label to avoid visual noise. | |||
| ### User Price Comparison | |||
| When a model has `display_pricing`, hide the old `UserPriceComparison` discount card. The detail side sheet should show the display pricing table and the estimate note instead. The current discount card is derived from real billing fields and would be misleading for models whose visible price is only an estimate. | |||
| `/api/pricing/user/*model` may still return `display_pricing` for consistency, but the frontend should not calculate user-specific discounted display pricing from `original_price` or `user_price`. | |||
| ### Admin Editing UI | |||
| The first implementation may manage `ModelDisplayPricing` through JSON option editing or a minimal admin form. If a structured form is added, it should expose only display fields: | |||
| - `specification` | |||
| - `official_supplier_tip` | |||
| - `price` as USD base price | |||
| - `unit` | |||
| - `sort_order` | |||
| - `discount_rate` | |||
| The editor must label `price` as USD base display price and include a warning that it does not affect real billing. | |||
| ## Billing Isolation | |||
| The implementation must not reference display pricing in: | |||
| - `relay/helper/price.go` | |||
| - `relay/helper/pricing_lookup.go` | |||
| - `controller/relay.go` pricing preload logic | |||
| - `service/task_polling.go` | |||
| - `types.PricingDecision` | |||
| - task final deduction calculation | |||
| - quota pre-consumption | |||
| - consume logs used for settlement | |||
| Only admin display-pricing config APIs, `model.Pricing`, and pricing-page APIs should expose it. Relay, billing, task, and consume-log paths must not read it. | |||
| ## Data Flow | |||
| ```text | |||
| Admin saves display pricing | |||
| -> options.ModelDisplayPricing | |||
| -> ratio_setting display pricing cache | |||
| -> model.RefreshPricing() | |||
| -> model.updatePricing() attaches display_pricing | |||
| -> GET /api/pricing returns display_pricing | |||
| -> frontend price page renders display SKU table | |||
| User creates a task | |||
| -> relay pricing preload reads ModelPricingRules or legacy price/ratio | |||
| -> PricingDecision is frozen | |||
| -> task completes | |||
| -> final settlement uses real billing data | |||
| -> display_pricing is never read | |||
| ``` | |||
| ## Testing | |||
| Backend tests: | |||
| - Parse and validate `ModelDisplayPricing` JSON. | |||
| - Default `discount_rate` to `1`. | |||
| - Reject invalid model keys, empty specification, negative price, empty unit, and invalid discount rates. | |||
| - Confirm `model.InitOptionMap()` initializes `ModelDisplayPricing` and `model.updateOptionMap()` hydrates the display pricing cache from DB-loaded option values. | |||
| - Confirm `model.GetPricingByModel()` includes sorted `DisplayPricing`. | |||
| - Confirm real `PricingConfig` behavior is unchanged when display pricing exists. | |||
| - Confirm `/api/pricing/user/*model` returns `display_pricing` but does not use it for computed prices. | |||
| - Confirm wildcard admin routes support model names containing `/`. | |||
| Frontend tests: | |||
| - `hasDisplayPricing()` detects display pricing before matrix pricing. | |||
| - Summary renders min-max range with converted currency. | |||
| - Detail table renders specification, supplier tip, unit, and discount. | |||
| - Existing matrix and legacy pricing rendering remains unchanged when no display pricing exists. | |||
| - User discount comparison is hidden when display pricing exists. | |||
| Manual verification: | |||
| - Configure one model with display pricing and token-based real billing. | |||
| - Confirm public price page shows the display table. | |||
| - Run a real task and confirm quota deduction follows the existing billing path, not `display_pricing.price`. | |||
| @@ -0,0 +1,184 @@ | |||
| # 天翼云 Seedance 渠道设计文档 | |||
| **日期**: 2026-07-08 | |||
| **状态**: 待用户审阅 | |||
| **方案**: A - 新增独立渠道类型 `DoubaoVideoCompatibleTianyiYun` | |||
| --- | |||
| ## 1. 背景与目标 | |||
| 天翼云 Seedance 渠道提供与 Doubao/Seedance 类似的异步视频生成能力。截图中的上游 API 使用: | |||
| - 提交任务: `POST https://ai.ctaigw.cn/v1/contents/generations/tasks` | |||
| - 查询任务: `GET https://ai.ctaigw.cn/v1/contents/generations/tasks/{task_id}` | |||
| - 鉴权: `Authorization: Bearer {api_key}` | |||
| - 模型: | |||
| - `cdance2.0-0611` | |||
| - `cdance2.0-fast-0611` | |||
| 目标是在 new-api 中新增独立渠道类型,复用现有 Seedance/Doubao 对外入口与任务链路,不把天翼云行为混入火山 Doubao 或 Aiping 渠道。 | |||
| --- | |||
| ## 2. 范围 | |||
| ### 2.1 包含 | |||
| - 新增 channel type: `DoubaoVideoCompatibleTianyiYun` | |||
| - 默认 base URL: `https://ai.ctaigw.cn` | |||
| - 复用现有 new-api 侧 Seedance/Doubao 入口: | |||
| - `POST /api/v3/contents/generations/tasks` | |||
| - `GET /api/v3/contents/generations/tasks/{task_id}` | |||
| - 新增任务适配器,转发到天翼云上游: | |||
| - `POST /v1/contents/generations/tasks` | |||
| - `GET /v1/contents/generations/tasks/{task_id}` | |||
| - 支持模型列表: | |||
| - `cdance2.0-0611` | |||
| - `cdance2.0-fast-0611` | |||
| - 前端渠道类型枚举增加新渠道,管理员可在渠道管理中创建该渠道。 | |||
| - 实现后使用用户提供的 key 发起一次真实最短任务验证。 | |||
| ### 2.2 不包含 | |||
| - 不新增用户侧公开路由 `/v1/contents/generations/tasks` | |||
| - 不修改现有火山 Doubao 官方渠道语义 | |||
| - 不修改现有 `DoubaoVideoCompatibleAiping` 渠道语义 | |||
| - 不在代码、文档、日志或测试 fixture 中保存用户提供的真实 API key | |||
| --- | |||
| ## 3. 总体设计 | |||
| 新增渠道保持独立: | |||
| ```text | |||
| 用户请求 | |||
| -> POST /api/v3/contents/generations/tasks | |||
| -> 现有 token auth / distribute / task relay | |||
| -> ChannelTypeDoubaoVideoCompatibleTianyiYun | |||
| -> doubao_tianyiyun TaskAdaptor | |||
| -> POST https://ai.ctaigw.cn/v1/contents/generations/tasks | |||
| ``` | |||
| 查询任务: | |||
| ```text | |||
| 用户请求 | |||
| -> GET /api/v3/contents/generations/tasks/{public_task_id} | |||
| -> 本地任务表定位 upstream task id | |||
| -> doubao_tianyiyun FetchTask | |||
| -> GET https://ai.ctaigw.cn/v1/contents/generations/tasks/{upstream_task_id} | |||
| ``` | |||
| 这样做的理由: | |||
| - 渠道类型边界清楚,排查日志、计费、状态映射时不会混淆 Aiping 与天翼云。 | |||
| - 对外入口复用已有 Seedance/Doubao 任务接口,客户端不需要学习新的 new-api 路由。 | |||
| - 上游路径差异只封装在新适配器中,后续如果天翼云返回结构变化,影响范围较小。 | |||
| --- | |||
| ## 4. 模型与请求体 | |||
| 用户侧模型名直接使用天翼云模型 ID: | |||
| - `cdance2.0-0611` | |||
| - `cdance2.0-fast-0611` | |||
| 管理员仍可使用现有模型映射能力,把内部展示模型映射到天翼云模型 ID。 | |||
| 请求体字段优先沿用现有 Doubao/Aiping Seedance 结构: | |||
| - `model` | |||
| - `content` | |||
| - `ratio` | |||
| - `duration` | |||
| - `watermark` | |||
| - 其它已由 Doubao/Aiping 适配器支持的 Seedance 字段,在不改变语义的情况下尽量透传。 | |||
| 适配器不额外把 `duration` 改成 `seconds`,因为截图示例明确使用 `duration`。 | |||
| --- | |||
| ## 5. 状态与响应处理 | |||
| 提交任务时,上游返回的 task id 存入本地任务记录;返回给用户的仍是 new-api public task id,避免暴露不同上游 ID 体系。 | |||
| 查询任务时,状态映射采用现有异步任务语义: | |||
| | 上游状态 | new-api 状态 | | |||
| | --- | --- | | |||
| | `pending`, `queued` | queued | | |||
| | `processing`, `running` | in_progress | | |||
| | `succeeded`, `success` | success | | |||
| | `failed`, `expired`, `cancelled` | failure | | |||
| | 其它未知状态 | in_progress | | |||
| 成功时从上游响应中提取视频 URL。优先兼容现有字段 `content.video_url`;如果真实验证发现字段不同,再按真实返回补充解析。 | |||
| --- | |||
| ## 6. 渠道测试与真实验证 | |||
| 普通渠道“测试”按钮不适合直接测试异步视频渠道。该新渠道应与现有 Doubao/Vidu 类似,避免用聊天补全接口做快速测试。 | |||
| 实现后的验证分两层: | |||
| 1. 自动化测试 | |||
| - 单元测试覆盖上游提交 URL、查询 URL、Authorization header、模型列表。 | |||
| - mock 上游测试覆盖提交响应 task id 与查询状态解析。 | |||
| - 回归测试确认现有 Doubao/Aiping 测试不受影响。 | |||
| 2. 真实上游验证 | |||
| - 使用用户本轮提供的 API key 作为运行时输入,不写入仓库。 | |||
| - 发起一次最短 5 秒任务,使用 `cdance2.0-0611` 或 `cdance2.0-fast-0611`。 | |||
| - 记录状态码、public task id、upstream task id、查询状态和必要错误信息。 | |||
| - 输出中不回显 API key。 | |||
| --- | |||
| ## 7. 需要修改的模块 | |||
| 预期实现会触及: | |||
| - `constant/channel.go` | |||
| - 在 `ChannelTypeDummy` 前新增渠道类型。 | |||
| - 同步 `ChannelBaseURLs` 与 `ChannelTypeNames`。 | |||
| - `relay/channel/task/doubao_tianyiyun/` | |||
| - 新增 `constants.go` | |||
| - 新增 `adaptor.go` | |||
| - 新增测试文件。 | |||
| - `relay/relay_adaptor.go` | |||
| - 注册新 task adaptor。 | |||
| - `common/api_type.go` | |||
| - 将新渠道映射到合适的 API type,建议沿用 VolcEngine/Doubao 视频相关路径。 | |||
| - `common/endpoint_type.go` | |||
| - 新渠道支持 `EndpointTypeDoubaoVideo`。 | |||
| - `controller/channel-test.go` | |||
| - 将新异步视频渠道加入不支持普通快速测试的列表。 | |||
| - `web/src/constants/channel.constants.js` | |||
| - 增加新渠道枚举,方便后台创建渠道。 | |||
| 根据实现中发现的实际依赖,可能还需要补充与模型能力、计费矩阵相关的白名单,但应保持最小改动。 | |||
| --- | |||
| ## 8. 风险与约束 | |||
| - `ChannelBaseURLs` 是按 channel type 索引的数组,新增类型必须同步维护索引,否则会造成渠道默认 base URL 错位。 | |||
| - 当前工作区已有其它未提交改动,实现时必须只修改本设计相关文件,不回退用户或其它任务的改动。 | |||
| - 真实验证会消耗上游额度;只做一次最短任务验证,除非用户要求重复验证。 | |||
| - 如果天翼云真实返回结构与截图示例不一致,应以真实响应为准补充解析,但不要扩大到无关功能。 | |||
| --- | |||
| ## 9. 验收标准 | |||
| - 管理后台能选择并创建 `DoubaoVideoCompatibleTianyiYun` 渠道。 | |||
| - 渠道默认 base URL 为 `https://ai.ctaigw.cn`。 | |||
| - 模型列表包含 `cdance2.0-0611` 与 `cdance2.0-fast-0611`。 | |||
| - 经由现有 `/api/v3/contents/generations/tasks` 提交时,上游请求落到 `/v1/contents/generations/tasks`。 | |||
| - 查询任务时,上游请求落到 `/v1/contents/generations/tasks/{task_id}`。 | |||
| - 单元测试和相关回归测试通过。 | |||
| - 使用真实 key 完成一次任务提交验证,并报告 task id 或明确的上游错误。 | |||
| @@ -0,0 +1,197 @@ | |||
| # 视频多维定价表 — 手动验证清单 | |||
| > 基于 `docs/testing/video-pricing-test-plan.md` | 测试环境:WSL CN (23000) | |||
| ## 环境准备 | |||
| ```bash | |||
| BASE=http://127.0.0.1:23000 | |||
| ADMIN_COOKIE="session_m=<从浏览器获取>" | |||
| TOKEN="sk-<admin-token>" | |||
| ``` | |||
| 登录方式:打开 `$BASE` → 登录 → DevTools → Application → Cookies → 复制 `session_m` 的值。 | |||
| --- | |||
| ## 一、CRUD 操作(P01–P12) | |||
| ### P01 获取配置列表 | |||
| ```bash | |||
| curl -s "$BASE/api/option/model_pricing" \ | |||
| -H "Cookie: $ADMIN_COOKIE" -H "New-Api-User: 1" | jq | |||
| ``` | |||
| ### P02 创建定价配置 | |||
| ```bash | |||
| curl -s -X PUT "$BASE/api/option/model_pricing/hailuo-video" \ | |||
| -H "Cookie: $ADMIN_COOKIE" -H "New-Api-User: 1" \ | |||
| -H "Content-Type: application/json" -d '{ | |||
| "schema_version": 1, "scope": "model", | |||
| "billing_unit": "per_call", "preconsume_strategy": "exact", | |||
| "dimensions": [ | |||
| {"key": "resolution", "source": "request.resolution", "type": "string", "default": null}, | |||
| {"key": "duration", "source": "request.duration", "type": "number", "default": null} | |||
| ], | |||
| "table": [ | |||
| {"resolution": "768P", "duration": 10, "price": 0.04, "source": "manual"}, | |||
| {"resolution": "1080P", "duration": 10, "price": 0.08, "source": "manual"} | |||
| ], | |||
| "fallback": {"strategy": "reject"} | |||
| }' | jq | |||
| ``` | |||
| ### P03 读取单配置 | |||
| ```bash | |||
| curl -s "$BASE/api/option/model_pricing/hailuo-video" \ | |||
| -H "Cookie: $ADMIN_COOKIE" -H "New-Api-User: 1" | jq | |||
| ``` | |||
| ### P05 删除配置 | |||
| ```bash | |||
| curl -s -X DELETE "$BASE/api/option/model_pricing/hailuo-video" \ | |||
| -H "Cookie: $ADMIN_COOKIE" -H "New-Api-User: 1" | jq | |||
| ``` | |||
| --- | |||
| ## 二、配置校验(V01–V17) | |||
| 使用 P02 的 body 修改字段来触发各种校验错误: | |||
| ### V07 per_call + minimum 冲突 | |||
| ```bash | |||
| curl -s -X PUT "$BASE/api/option/model_pricing/test-v" \ | |||
| -H "Cookie: $ADMIN_COOKIE" -H "New-Api-User: 1" \ | |||
| -H "Content-Type: application/json" -d '{ | |||
| "schema_version": 1, "scope": "model", | |||
| "billing_unit": "per_call", "preconsume_strategy": "minimum", | |||
| "dimensions": [{"key":"r","source":"request.resolution","type":"string","default":null}], | |||
| "table": [{"r":"x","price":0.04,"source":"manual"}], | |||
| "fallback": {"strategy": "reject"} | |||
| }' | jq | |||
| # 预期: {"success": false, "message": "billing_unit per_call requires..."} | |||
| ``` | |||
| --- | |||
| ## 三、计费验证(E01–E08) | |||
| 需要先配置定价,再提交视频任务,观察 quota 变化。 | |||
| ### E01 精确匹配 — 提交任务 | |||
| ```bash | |||
| # 1. 配置定价表 | |||
| curl -s -X PUT "$BASE/api/option/model_pricing/hailuo-video" \ | |||
| -H "Cookie: $ADMIN_COOKIE" -H "New-Api-User: 1" \ | |||
| -H "Content-Type: application/json" -d '{ | |||
| "schema_version": 1, "scope": "model", | |||
| "billing_unit": "per_call", "preconsume_strategy": "exact", | |||
| "dimensions": [ | |||
| {"key": "res", "source": "request.resolution", "type": "string", "default": null} | |||
| ], | |||
| "table": [{"res": "768P", "price": 0.04, "source": "manual"}], | |||
| "fallback": {"strategy": "reject"} | |||
| }' | |||
| # 2. 查询当前配额 | |||
| Q_BEFORE=$(curl -s "$BASE/api/user/self" \ | |||
| -H "Authorization: Bearer $TOKEN" | jq '.data.quota') | |||
| echo "Quota before: $Q_BEFORE" | |||
| # 3. 提交视频任务 | |||
| TASK_ID=$(curl -s -X POST "$BASE/v1/videos" \ | |||
| -H "Authorization: Bearer $TOKEN" \ | |||
| -H "Content-Type: application/json" \ | |||
| -d '{"model":"hailuo-video","prompt":"test","resolution":"768P"}' \ | |||
| | jq -r '.id') | |||
| echo "Task: $TASK_ID" | |||
| # 4. 等待完成(约 45 秒) | |||
| sleep 50 | |||
| # 5. 查询任务结果 | |||
| curl -s "$BASE/v1/videos/$TASK_ID" \ | |||
| -H "Authorization: Bearer $TOKEN" | jq | |||
| # 6. 查询配额 | |||
| Q_AFTER=$(curl -s "$BASE/api/user/self" \ | |||
| -H "Authorization: Bearer $TOKEN" | jq '.data.quota') | |||
| echo "Quota after: $Q_AFTER" | |||
| echo "Delta: $((Q_BEFORE - Q_AFTER))" | |||
| # 预期 delta = 20000 (= 0.04 × 500000) | |||
| ``` | |||
| ### E04 不匹配 → reject | |||
| ```bash | |||
| curl -i -X POST "$BASE/v1/videos" \ | |||
| -H "Authorization: Bearer $TOKEN" \ | |||
| -H "Content-Type: application/json" \ | |||
| -d '{"model":"hailuo-video","prompt":"reject test","resolution":"NOT_IN_TABLE"}' | |||
| # 预期: HTTP 400, quota 不变 | |||
| ``` | |||
| ### E03 Wildcard 匹配 | |||
| 配置 `"res": "*"` + `"dur": 10`,发送任意 resolution → 匹配 wildcard 行 → quota 扣 0.05×500000=25000 | |||
| ### E05 Fallback=max | |||
| 配置 `"fallback": {"strategy": "max"}`,发不匹配请求 → 按表中最高价 0.10 计费 | |||
| ### E06 Fallback=default | |||
| 配置 `"fallback": {"strategy": "default", "default_price": 0.07}` → 不匹配时按 0.07×500000=35000 计费 | |||
| --- | |||
| ## 四、Per 1M Tokens 计费(E09–E14,需 mock 上游) | |||
| ```bash | |||
| # 配置 Doubao/Seedance 的 per_1m_tokens 定价 | |||
| curl -s -X PUT "$BASE/api/option/model_pricing/seedance-2" \ | |||
| -H "Cookie: $ADMIN_COOKIE" -H "New-Api-User: 1" \ | |||
| -H "Content-Type: application/json" -d '{ | |||
| "schema_version": 1, "scope": "model", | |||
| "billing_unit": "per_1m_tokens", "preconsume_strategy": "minimum", | |||
| "dimensions": [ | |||
| {"key": "res", "source": "request.resolution", "type": "string", "default": null} | |||
| ], | |||
| "table": [{"res": "1080P", "price": 200, "source": "manual"}], | |||
| "fallback": {"strategy": "reject"} | |||
| }' | |||
| # 提交任务(需 DoubaoVideo 渠道) | |||
| curl -s -X POST "$BASE/v1/videos" \ | |||
| -H "Authorization: Bearer $TOKEN" \ | |||
| -H "Content-Type: application/json" \ | |||
| -d '{"model":"seedance-2","prompt":"test","resolution":"1080P","_mock_total_tokens":1000000}' | |||
| # 预扣: PreConsumedQuota × 200 × ratio / 1M ≈ 1 quota | |||
| # 结算: 1000000 × 200 / 1M × QuotaPerUnit × ratio | |||
| ``` | |||
| --- | |||
| ## 五、Remix 计费(E15–E18) | |||
| ```bash | |||
| # 对已完成的视频任务进行 remix | |||
| curl -s -X POST "$BASE/v1/videos/$ORIGINAL_VIDEO_ID/remix" \ | |||
| -H "Authorization: Bearer $TOKEN" \ | |||
| -H "Content-Type: application/json" \ | |||
| -d '{"model":"hailuo-video","prompt":"remix version"}' | |||
| # 预期: 继承原任务的定价配置, 预扣与原始任务相同 | |||
| ``` | |||
| --- | |||
| ## 快速自检清单 | |||
| | # | 操作 | 验证点 | ✓ | | |||
| |---|------|--------|---| | |||
| | 1 | 创建配置 | API 返回 200,含完整配置 | | | |||
| | 2 | 非法 schema_version=2 | API 返回 400 | | | |||
| | 3 | 精确匹配提交任务 | quota 扣除 = price × 500000 | | | |||
| | 4 | 不匹配+reject | 400 + quota 不变 | | | |||
| | 5 | Wildcard 匹配 | 正确匹配 * 行 | | | |||
| | 6 | Fallback max | 扣表中最高价 | | | |||
| | 7 | Fallback default | 扣 default_price | | | |||
| | 8 | 删除配置 | 200 | | | |||
| | 9 | 删除后读取 | 404 | | | |||
| @@ -0,0 +1,198 @@ | |||
| # Video Pricing Table Testing | |||
| This guide verifies task video billing with model-level pricing tables. | |||
| ## Prerequisites | |||
| - Admin session cookie for `/api/option/model_pricing`. | |||
| - API token for `/v1/videos`. | |||
| - A usable Hailuo channel for `hailuo-video`. | |||
| - A usable Seedance/Doubao channel for `seedance-2`; for deterministic usage billing, point it at a mock upstream that returns `usage.total_tokens`. | |||
| - `QuotaPerUnit = 500000` unless your environment overrides it. | |||
| Set shell variables: | |||
| ```bash | |||
| BASE_URL=http://127.0.0.1:3000 | |||
| ADMIN_COOKIE='session=...' | |||
| API_TOKEN='sk-...' | |||
| ``` | |||
| ## Hailuo Per Call | |||
| Configure an exact per-call matrix: | |||
| ```bash | |||
| curl -sS -X PUT "$BASE_URL/api/option/model_pricing/hailuo-video" \ | |||
| -H "Cookie: $ADMIN_COOKIE" \ | |||
| -H 'Content-Type: application/json' \ | |||
| -d '{ | |||
| "schema_version": 1, | |||
| "scope": "model", | |||
| "billing_unit": "per_call", | |||
| "preconsume_strategy": "exact", | |||
| "dimensions": [ | |||
| {"key": "resolution", "source": "request.resolution", "type": "string", "default": null}, | |||
| {"key": "duration", "source": "request.duration", "type": "number", "default": null} | |||
| ], | |||
| "table": [ | |||
| {"resolution": "768P", "duration": 10, "price": 0.04, "source": "manual"} | |||
| ], | |||
| "fallback": {"strategy": "reject"} | |||
| }' | |||
| ``` | |||
| Submit a matched request: | |||
| ```bash | |||
| curl -sS -X POST "$BASE_URL/v1/videos" \ | |||
| -H "Authorization: Bearer $API_TOKEN" \ | |||
| -H 'Content-Type: application/json' \ | |||
| -d '{ | |||
| "model": "hailuo-video", | |||
| "prompt": "a pricing smoke test", | |||
| "resolution": "768P", | |||
| "duration": 10 | |||
| }' | |||
| ``` | |||
| Expected quota: | |||
| ```text | |||
| actual_quota = price_usd * QuotaPerUnit * group_ratio | |||
| = 0.04 * 500000 * 1 | |||
| = 20000 | |||
| ``` | |||
| Submit an unsupported row and expect HTTP 400 with no quota change: | |||
| ```bash | |||
| curl -i -X POST "$BASE_URL/v1/videos" \ | |||
| -H "Authorization: Bearer $API_TOKEN" \ | |||
| -H 'Content-Type: application/json' \ | |||
| -d '{ | |||
| "model": "hailuo-video", | |||
| "prompt": "a pricing reject test", | |||
| "resolution": "1080P", | |||
| "duration": 10 | |||
| }' | |||
| ``` | |||
| ## Seedance Per 1M Tokens | |||
| Configure a usage matrix: | |||
| ```bash | |||
| curl -sS -X PUT "$BASE_URL/api/option/model_pricing/seedance-2" \ | |||
| -H "Cookie: $ADMIN_COOKIE" \ | |||
| -H 'Content-Type: application/json' \ | |||
| -d '{ | |||
| "schema_version": 1, | |||
| "scope": "model", | |||
| "billing_unit": "per_1m_tokens", | |||
| "preconsume_strategy": "minimum", | |||
| "dimensions": [ | |||
| {"key": "resolution", "source": "request.resolution", "type": "string", "default": null} | |||
| ], | |||
| "table": [ | |||
| {"resolution": "720P", "price": 0.5, "source": "manual"} | |||
| ], | |||
| "fallback": {"strategy": "reject"} | |||
| }' | |||
| ``` | |||
| Submit: | |||
| ```bash | |||
| curl -sS -X POST "$BASE_URL/v1/videos" \ | |||
| -H "Authorization: Bearer $API_TOKEN" \ | |||
| -H 'Content-Type: application/json' \ | |||
| -d '{ | |||
| "model": "seedance-2", | |||
| "prompt": "a usage pricing smoke test", | |||
| "resolution": "720P", | |||
| "duration": 5 | |||
| }' | |||
| ``` | |||
| Fetch until complete: | |||
| ```bash | |||
| TASK_ID=task_xxx | |||
| curl -sS "$BASE_URL/v1/videos/$TASK_ID" \ | |||
| -H "Authorization: Bearer $API_TOKEN" | |||
| ``` | |||
| When the upstream completion response contains `usage.total_tokens = 1000`, expected final quota is: | |||
| ```text | |||
| actual_quota = total_tokens * token_unit_price_usd / 1_000_000 * QuotaPerUnit * group_ratio | |||
| = 1000 * 0.5 / 1_000_000 * 500000 * 1 | |||
| = 250 | |||
| ``` | |||
| The settlement must use the frozen `token_unit_price_usd` stored on the task billing context. Changing or deleting the current `ModelPricingRules` option after submit must not affect the completion calculation for that task. | |||
| ### Seedance Video Input Pricing | |||
| Seedance 2.0 can use `derived.video_input` to select a different token price when the request includes a video input. This replaces hard-coded video-input discounts with a configurable pricing table. | |||
| For `doubao-seedance-2-0-260128`: | |||
| ```json | |||
| { | |||
| "schema_version": 1, | |||
| "scope": "model", | |||
| "billing_unit": "per_1m_tokens", | |||
| "preconsume_strategy": "minimum", | |||
| "dimensions": [ | |||
| { "key": "video_input", "source": "derived.video_input", "type": "boolean", "default": null } | |||
| ], | |||
| "table": [ | |||
| { "video_input": false, "price": 46, "source": "manual" }, | |||
| { "video_input": true, "price": 28, "source": "manual" } | |||
| ], | |||
| "fallback": { "strategy": "reject" } | |||
| } | |||
| ``` | |||
| For `doubao-seedance-2-0-fast-260128`, use the same dimension and fallback with these rows: | |||
| ```json | |||
| [ | |||
| { "video_input": false, "price": 37, "source": "manual" }, | |||
| { "video_input": true, "price": 22, "source": "manual" } | |||
| ] | |||
| ``` | |||
| `derived.video_input` is true when the request has a top-level video field such as `video_url`, or when `metadata.content` contains a `video_url` item. | |||
| ## Delete Config And Legacy Fallback | |||
| Delete a matrix config: | |||
| ```bash | |||
| curl -sS -X DELETE "$BASE_URL/api/option/model_pricing/hailuo-video" \ | |||
| -H "Cookie: $ADMIN_COOKIE" | |||
| curl -sS -X DELETE "$BASE_URL/api/option/model_pricing/seedance-2" \ | |||
| -H "Cookie: $ADMIN_COOKIE" | |||
| ``` | |||
| After deletion, requests for that model return to the legacy `ModelPrice` path. The `ModelPrice` option itself is not removed or modified by the model pricing table API. | |||
| ## Smoke Script | |||
| The same flow can be run with: | |||
| ```bash | |||
| "C:/Users/28221/.conda/envs/py312/python.exe" test-scripts/test_video_pricing.py | |||
| ``` | |||
| Required environment: | |||
| ```bash | |||
| NEW_API_BASE_URL=http://127.0.0.1:3000 | |||
| NEW_API_ADMIN_USERNAME=root | |||
| NEW_API_ADMIN_PASSWORD=123456 | |||
| NEW_API_TOKEN=sk-... | |||
| ``` | |||
| @@ -0,0 +1,341 @@ | |||
| # 视频定价表测试方案 | |||
| > 分支:`feat/video-pricing-table-codex` | 日期:2026-06-18 | |||
| ## 概述 | |||
| 本文档定义了多维定价表系统的完整测试方案,包括两部分: | |||
| - **手动验证清单**:覆盖全部功能点,含 curl 示例和预期结果 | |||
| - **自动化测试脚本**(Python):覆盖核心 CRUD + 计费路径 | |||
| ## 测试环境 | |||
| ```bash | |||
| BASE_URL=http://127.0.0.1:3000 # 或 http://127.0.0.1:23000 (WSL CN) | |||
| ADMIN_USERNAME=root | |||
| ADMIN_PASSWORD=123456 | |||
| API_TOKEN=sk-... # 任意有效用户 token,配额充足 ≥ 100000 | |||
| ``` | |||
| 前置条件:服务运行中、测试用户存在、Hailuo/Doubao 渠道可用、`QuotaPerUnit = 500000`。 | |||
| --- | |||
| ## 一、定价配置 CRUD(P01–P12) | |||
| | 编号 | 测试点 | 方法 | 路径 | 预期 | | |||
| |------|--------|------|------|------| | |||
| | P01 | 获取空配置列表 | GET | `/api/option/model_pricing` | 200, `{}` | | |||
| | P02 | 创建配置 | PUT | `/api/option/model_pricing/hailuo-video` | 200, 返回配置 | | |||
| | P03 | 读取单配置 | GET | `/api/option/model_pricing/hailuo-video` | 200 | | |||
| | P04 | 更新配置 | PUT | 同上 + 修改后 table | 200 | | |||
| | P05 | 删除配置 | DELETE | `/api/option/model_pricing/hailuo-video` | 200 | | |||
| | P06 | 删除后读取 | GET | 同上 | 404 | | |||
| | P07 | 批量列表 | GET | `/api/option/model_pricing` | 200, 含所有已创建 | | |||
| | P08 | 未认证请求 | PUT | 不带 Cookie | 401/403 | | |||
| | P09 | 错误 Content-Type | PUT | `Content-Type: text/plain` | 400 | | |||
| | P10 | 非法 JSON | PUT | 损坏 JSON | 400 | | |||
| | P11 | 幂等覆盖 | PUT | 再次 PUT 已存在 | 200 | | |||
| | P12 | 删除不存在 | DELETE | `.../nonexistent` | 200 | | |||
| ### P02 请求体示例 | |||
| ```json | |||
| { | |||
| "schema_version": 1, | |||
| "scope": "model", | |||
| "billing_unit": "per_call", | |||
| "preconsume_strategy": "exact", | |||
| "dimensions": [ | |||
| {"key": "resolution", "source": "request.resolution", "type": "string", "default": null}, | |||
| {"key": "duration", "source": "request.duration", "type": "number", "default": null} | |||
| ], | |||
| "table": [ | |||
| {"resolution": "768P", "duration": 10, "price": 0.04, "source": "manual"} | |||
| ], | |||
| "fallback": {"strategy": "reject"} | |||
| } | |||
| ``` | |||
| --- | |||
| ## 二、配置校验规则(V01–V17) | |||
| | 编号 | 校验规则 | 触发方式 | 预期 | | |||
| |------|---------|---------|------| | |||
| | V01 | schema_version ≠ 1 | `"schema_version": 2` | 400 | | |||
| | V02 | scope ≠ model | `"scope": "channel"` | 400 | | |||
| | V03 | dimensions 为空 | `"dimensions": []` | 400 | | |||
| | V04 | table 为空 | `"table": []` | 400 | | |||
| | V05 | billing_unit 非法 | `"billing_unit": "per_hour"` | 400 | | |||
| | V06 | preconsume_strategy 非法 | `"preconsume_strategy": "average"` | 400 | | |||
| | V07 | per_call + minimum 冲突 | billing_unit=per_call, strategy=minimum | 400 | | |||
| | V08 | per_1m_tokens + exact 冲突 | billing_unit=per_1m_tokens, strategy=exact | 400 | | |||
| | V09 | 不支持 matrix usage 模型 | hailuo-video + per_1m_tokens | 400 | | |||
| | V10 | dimension key 保留字 | key="price" | 400 | | |||
| | V11 | dimension key 非法格式 | key="reso lution" | 400 | | |||
| | V12 | dimension type 非法 | type="integer" | 400 | | |||
| | V13 | dimension source 非法 | source="header.X-Custom" | 400 | | |||
| | V14 | row 缺少维度 | table 行不含某维度值 | 400 | | |||
| | V15 | row price 负数 | `"price": -1` | 400 | | |||
| | V16 | 模糊冲突行 | 相同 wildcard count 的行重叠 | 400 | | |||
| | V17 | optional 缺 default | optional=true + default=null | 400 | | |||
| --- | |||
| ## 三、维度解析(D01–D06) | |||
| | 编号 | 测试点 | 说明 | | |||
| |------|--------|------| | |||
| | D01 | JSON 请求体浅层提取 | `{"resolution": "1080P"}` → 字符串 | | |||
| | D02 | JSON 嵌套提取 | `request.options.size` 深层路径 | | |||
| | D03 | Multipart 提取 | `Content-Type: multipart/form-data` | | |||
| | D04 | Optional 默认值 | 请求不传 → 用 default | | |||
| | D05 | 缺失必填维度 | 该维度为 nil,不匹配任何行 | | |||
| | D06 | Derived 维度 | megapixels/aspect_ratio/input_mode 自动计算 | | |||
| --- | |||
| ## 四、定价匹配 + 回退策略(L01–L08) | |||
| | 编号 | 场景 | 预期 | | |||
| |------|------|------| | |||
| | L01 | 精确匹配 | 维度 = 行值 → 返回该行价格 | | |||
| | L02 | Wildcard 匹配 | `*` 匹配任意值 | | |||
| | L03 | 最小 wildcard 优先 | 多行匹配选 wildcard 最少 | | |||
| | L04 | 无匹配 → reject | 400 + 不扣 quota | | |||
| | L05 | 无匹配 → max | 取表中最高价 | | |||
| | L06 | 无匹配 → default | 取 default_price | | |||
| | L07 | 歧义行配置拒绝 | V16 | | |||
| | L08 | per_1m_tokens 预扣 | minimum 策略预扣 | | |||
| --- | |||
| ## 五、Task 计费生命周期(B01–B08) | |||
| | 编号 | 场景 | 说明 | | |||
| |------|------|------| | |||
| | B01 | 提交 + 预扣 | PreConsume → DB quota 变化 | | |||
| | B02 | 成功 + 结算 | per_call 精确不补扣 | | |||
| | B03 | 失败 + 退款 | Refund 退还 | | |||
| | B04 | Adaptor 调整差额 | AdjustBillingOnSubmit → 差额 | | |||
| | B05 | 免费模型 | price=0 → 不预扣 | | |||
| | B06 | 分组倍率 | groupRatio≠1 → quota 按倍率 | | |||
| | B07 | BillingContext 持久化 | DB 记录完整 | | |||
| | B08 | 消费日志 | log 表字段正确 | | |||
| --- | |||
| ## 六、Matrix Usage 计费(M01–M06) | |||
| | 编号 | 场景 | 说明 | | |||
| |------|------|------| | |||
| | M01 | 预扣最小值 | ≈ PreConsumedQuota × tokenUnitPrice × ratio / 1M | | |||
| | M02 | Token 结算 | 上游返回 total_tokens → 按实际计算 | | |||
| | M03 | 补扣差额 | 实际 > 预扣 → 补扣 | | |||
| | M04 | 退还差额 | 实际 < 预扣 → 退还 | | |||
| | M05 | 不支持模型被拒 | hailuo-video + per_1m_tokens → 400 | | |||
| | M06 | 渠道类型不匹配 | 非 DoubaoVideo → 不支持 | | |||
| --- | |||
| ## 七、Remix 计费(R01–R05) | |||
| | 编号 | 场景 | 说明 | | |||
| |------|------|------| | |||
| | R01 | BillingContext 继承 | 复用原任务定价快照 | | |||
| | R02 | 旧格式兼容 | 从 seconds/size 解析 | | |||
| | R03 | 锁渠道 | 必须用原渠道 | | |||
| | R04 | 原任务不存在 | 400 | | |||
| | R05 | Matrix remix | 按相同方式计费 | | |||
| --- | |||
| ## 八、渠道选择(C01–C04) | |||
| | 编号 | 场景 | 说明 | | |||
| |------|------|------| | |||
| | C01 | Matrix 渠道路由 | 只选匹配渠道 | | |||
| | C02 | 无可用渠道 | 错误 | | |||
| | C03 | 普通模型不限渠道 | per_call 不需 matrix | | |||
| | C04 | 模型映射后匹配 | 上游映射后仍正确 | | |||
| --- | |||
| ## 九、实际计费请求(E01–E22) | |||
| ### Per Call 计费(E01–E08) | |||
| | 编号 | 场景 | 预期 | | |||
| |------|------|------| | |||
| | E01 | 精确匹配 768P/10s=0.04 | quota 扣 20000 | | |||
| | E02 | 提交前后 quota 差 | 差值 = price × QuotaPerUnit × groupRatio | | |||
| | E03 | Wildcard 匹配 | `*` 匹配任意同维度值 | | |||
| | E04 | 不匹配 → reject | quota 不变 + 400 | | |||
| | E05 | Fallback=max | 扣最高价 | | |||
| | E06 | Fallback=default | 扣 default_price × 500000 | | |||
| | E07 | 三维度匹配 | resolution × duration × input_mode | | |||
| | E08 | Optional 默认值 | 不传 → 自动补 default → 匹配 | | |||
| ### Per 1M Tokens 计费(E09–E14,需 mock 上游) | |||
| | 编号 | 场景 | 说明 | | |||
| |------|------|------| | |||
| | E09 | 预扣最小值 | 预扣 ≈ 1–2 quota | | |||
| | E10 | 按 tokens 结算 | tokens × price × ratio / 1M | | |||
| | E11 | 补扣差额 | tokens 多 → 补扣 | | |||
| | E12 | 退还差额 | tokens 少 → 退还 | | |||
| | E13 | 定价表 + usage | 维度匹配得单价 → × total_tokens | | |||
| | E14 | 预扣量对比 | per_1m_tokens vs per_call 预扣差 | | |||
| ### Remix 计费完整性(E15–E18) | |||
| | 编号 | 场景 | 说明 | | |||
| |------|------|------| | |||
| | E15 | Remix per_call | 扣费与原任务相同 | | |||
| | E16 | Remix per_1m_tokens | 相同 token 单价 | | |||
| | E17 | Remix 旧格式 | seconds/size 推断 | | |||
| | E18 | Remix 渠道不可用 | 400 | | |||
| ### 边界场景(E19–E22) | |||
| | 编号 | 场景 | 说明 | | |||
| |------|------|------| | |||
| | E19 | 定价配置优先于旧价格 | matrix 优先 | | |||
| | E20 | 无定价配置走旧逻辑 | legacy path | | |||
| | E21 | Quota 不足 | 提交失败 | | |||
| | E22 | GroupRatio 非 1 | quota 按倍率 | | |||
| --- | |||
| ## 自动化脚本范围 | |||
| `test-scripts/test_model_pricing_api.py`,预估覆盖约 33 个用例: | |||
| - **CRUD**:P01–P12(约 8 个) | |||
| - **校验规则**:V01–V17(约 12 个) | |||
| - **Per Call 计费**:E01–E08(约 6 个) | |||
| - **Per 1M Tokens**:E09–E14(约 4 个,需 mock 上游) | |||
| --- | |||
| ## 十、Mock 上游服务器 | |||
| 由于本地/测试环境没有真实的 Doubao 和 Hailuo 视频生成 API,需要一个 mock server 来模拟上游行为。 | |||
| ### 为什么需要 mock | |||
| | 测试类别 | 依赖上游做什么 | 没有 mock 的后果 | | |||
| |---------|--------------|-----------------| | |||
| | Per Call 计费(E01-E08) | 接受提交 → 后续 polling 返回 succeeded | 任务一直 pending,永远不结算 | | |||
| | Per 1M Tokens(E09-E14) | 提交 → polling 返回 succeeded + `usage.total_tokens` | 无法验证按 token 结算逻辑 | | |||
| | Remix(E15-E18) | 基于已完成任务再次提交 → polling 返回 succeeded | 无法验证 remix 计费链路 | | |||
| ### Mock 需要模拟的 API | |||
| #### Hailuo(Minimax)API | |||
| **提交任务** — `POST /v1/video/generation` | |||
| ```json | |||
| // Request (new-api 转发) | |||
| {"model": "hailuo-video", "prompt": "...", "resolution": "768P", "duration": 10} | |||
| // Response | |||
| {"data": {"task_id": "mock_hailuo_xxx"}, "base_resp": {"status_code": 0, "status_msg": "success"}} | |||
| ``` | |||
| **查询任务** — `GET /v1/query/video_result?task_id={id}` | |||
| ```json | |||
| // Response (succeeded) | |||
| {"data": { | |||
| "status": "Success", | |||
| "video_url": "https://mock.local/video.mp4" | |||
| }, "base_resp": {"status_code": 0}} | |||
| ``` | |||
| #### Doubao(Seedance)API | |||
| **提交任务** — `POST /api/v3/contents/generations/tasks` | |||
| ```json | |||
| // Request (new-api 转发) | |||
| {"model": "seedance-2", "content": [{"type": "text", "text": "..."}], "resolution": "...", "duration": 5} | |||
| // Response | |||
| {"id": "mock_doubao_xxx"} | |||
| ``` | |||
| **查询任务** — `GET /api/v3/contents/generations/tasks/{id}` | |||
| ```json | |||
| // Response (succeeded + usage) | |||
| { | |||
| "id": "mock_doubao_xxx", | |||
| "model": "seedance-2", | |||
| "status": "succeeded", | |||
| "content": {"video_url": "https://mock.local/video.mp4"}, | |||
| "usage": {"completion_tokens": 10000, "total_tokens": 10000} | |||
| } | |||
| ``` | |||
| ### Mock 核心设计 | |||
| **关键思路:** mock 将请求中的信息存储起来,在 polling 时原样返回给 new-api。`total_tokens` 通过请求体中的特殊字段控制。 | |||
| **状态机:** | |||
| ``` | |||
| submit → "queued" (首次 GET) | |||
| → "processing" (第二次 GET) | |||
| → "succeeded" (第三次 GET,返回 usage) | |||
| ``` | |||
| **控制 total_tokens 的方式:** 在提交请求中包含 `_mock_total_tokens` 字段,mock 提取后不传给上游逻辑,只在 fetch 时返回该值作为 `usage.total_tokens`。 | |||
| **Python 实现框架:** | |||
| ```python | |||
| # mock_server.py - 使用 Python stdlib http.server,无外部依赖 | |||
| # 用法: python test-scripts/mock_video_upstream.py --port 18999 | |||
| import json | |||
| from http.server import HTTPServer, BaseHTTPRequestHandler | |||
| class MockHandler(BaseHTTPRequestHandler): | |||
| tasks = {} # task_id → {request, poll_count, total_tokens} | |||
| def do_POST(self): | |||
| body = json.loads(self.rfile.read(int(self.headers['Content-Length']))) | |||
| task_id = f"mock_{len(self.tasks):04d}" | |||
| # 提取 _mock_total_tokens 并移除 | |||
| total_tokens = body.pop("_mock_total_tokens", None) | |||
| self.tasks[task_id] = { | |||
| "request": body, | |||
| "poll_count": 0, | |||
| "total_tokens": total_tokens or 10000, | |||
| } | |||
| # 根据 path 返回 Hailuo 或 Doubao 格式 | |||
| ... | |||
| def do_GET(self): | |||
| # 解析 task_id,增加 poll_count | |||
| # poll_count 0: queued, 1: processing, 2+: succeeded | |||
| ... | |||
| ``` | |||
| ### 测试流程(使用 mock) | |||
| ``` | |||
| 1. 启动 mock server: python test-scripts/mock_video_upstream.py --port 18999 | |||
| 2. 在 new-api 中创建指向 mock 的渠道: | |||
| - Hailuo 渠道: base_url = http://host.docker.internal:18999 | |||
| - Doubao 渠道: base_url = http://host.docker.internal:18999, type = DoubaoVideo | |||
| 3. 运行测试脚本: python test-scripts/test_model_pricing_api.py | |||
| 4. 停止 mock server: Ctrl+C | |||
| ``` | |||
| ### Mock Server 配置端点 | |||
| | 端点 | 方法 | 用途 | | |||
| |------|------|------| | |||
| | `/mock/reset` | POST | 清空所有任务状态 | | |||
| | `/mock/stats` | GET | 查看已完成/进行中的任务数 | | |||
| | `/mock/tasks` | GET | 列出所有任务及其状态 | | |||
| - **Remix**:E15–E18(约 3 个) | |||
| @@ -0,0 +1,264 @@ | |||
| # 多维计费系统 — 测试验证报告 | |||
| > 日期:2026-06-18 | 分支:`feat/video-pricing-table-codex` | 测试环境:WSL CN (localhost:23000) | |||
| ## 一、测试概览 | |||
| | 维度 | 用例数 | 通过 | 失败 | 通过率 | | |||
| |------|--------|------|------|--------| | |||
| | P: 定价配置 CRUD | 9 | 9 | 0 | 100% | | |||
| | V: 配置校验规则 | 16 | 16 | 0 | 100% | | |||
| | D: 维度解析 | 2 | 2 | 0 | 100% | | |||
| | W: Wildcard 优先级 | 1 | 1 | 0 | 100% | | |||
| | E: Per Call 计费 | 7 | 7 | 0 | 100% | | |||
| | E: Per 1M Tokens 计费 | 3 | 3 | 0 | 100% | | |||
| | DB: 数据持久化 | 1 | 1 | 0 | 100% | | |||
| | **合计** | **39** | **39** | **0** | **100%** | | |||
| ## 二、测试架构 | |||
| ``` | |||
| ┌─────────────────────────────────────────────────────────────────┐ | |||
| │ Windows 主机 │ | |||
| │ │ | |||
| │ ┌──────────────────────┐ ┌───────────────────────────────┐ │ | |||
| │ │ test_model_pricing │ │ mock_video_upstream.py │ │ | |||
| │ │ _api.py (39 tests) │ │ :18999 │ │ | |||
| │ │ │ │ │ │ | |||
| │ │ requests.Session │ │ Hailuo: /v1/video_generation │ │ | |||
| │ │ cookie auth (CRUD) │ │ Doubao: /api/v3/.../tasks │ │ | |||
| │ │ token auth (billing) │ │ /mock/config 控制 total_tokens│ │ | |||
| │ └──────┬───────────────┘ └────────────┬──────────────────┘ │ | |||
| │ │ :23000 │ :18999 │ | |||
| └─────────┼───────────────────────────────────┼─────────────────────┘ | |||
| │ │ | |||
| ┌─────────┼───────────────────────────────────┼─────────────────────┐ | |||
| │ WSL │ │ │ | |||
| │ ┌──────▼───────────────────────────────────▼───────────────────┐ │ | |||
| │ │ new-api (Docker, cn-new-api) │ │ | |||
| │ │ :23000 │ │ | |||
| │ │ │ │ | |||
| │ │ 渠道: mock-hl (type=35/MiniMax) → http://172.22.48.1:18999 │ │ | |||
| │ │ mock-db (type=54/Doubao) → http://172.22.48.1:18999 │ │ | |||
| │ │ │ │ | |||
| │ │ MySQL: channels + abilities + tokens + pricing_configs │ │ | |||
| │ │ Redis: session cache + rate limiter │ │ | |||
| │ └───────────────────────────────────────────────────────────────┘ │ | |||
| └───────────────────────────────────────────────────────────────────┘ | |||
| ``` | |||
| ### Mock Server 核心机制 | |||
| ``` | |||
| 状态机: POST submit → poll_count=0 → "queued" | |||
| GET fetch 1 → poll_count=1 → "processing" | |||
| GET fetch 2 → poll_count=2 → "processing" | |||
| GET fetch 3+ → poll_count≥3 → "succeeded" | |||
| 控制手段: POST /mock/config {"total_tokens": 500000} | |||
| → 创建任务时使用该值作为 usage.total_tokens | |||
| → per_1m_tokens 结算 = total_tokens × price / 1M × QuotaPerUnit | |||
| ``` | |||
| ## 三、逐模块验证结果 | |||
| ### 3.1 定价配置 CRUD(P01-P12) | |||
| | 用例 | 验证点 | 请求 | 结果 | | |||
| |------|--------|------|------| | |||
| | P01 | 空列表 | GET /api/option/model_pricing | 200, {} | | |||
| | P02 | 创建配置 | PUT .../test-model-1 + 合法 JSON | 200, 配置已存储 | | |||
| | P03 | 读取单配置 | GET .../test-model-1 | 200, scope=model | | |||
| | P04 | 更新配置 | PUT .../test-model-1 + 新 table | 200, table 已更新 | | |||
| | P05 | 删除配置 | DELETE .../test-model-1 | 200 | | |||
| | P06 | 删除后读取 | GET .../test-model-1 | 404 | | |||
| | P07 | 批量列表 | GET /api/option/model_pricing | 含所有已创建 | | |||
| | P11 | 幂等覆盖 | 两次 PUT 同 model | 200, 不报错 | | |||
| | P12 | 删除不存在 | DELETE .../no-such | 200 | | |||
| ### 3.2 配置校验规则(V01-V17) | |||
| 所有校验规则均正确拒绝非法输入,返回 400: | |||
| | 规则 | 触发条件 | 验证 | | |||
| |------|---------|------| | |||
| | V01 | schema_version ≠ 1 | ✅ | | |||
| | V02 | scope ≠ "model" | ✅ | | |||
| | V03 | dimensions = [] | ✅ | | |||
| | V04 | table = [] | ✅ | | |||
| | V05 | billing_unit 非法 | ✅ | | |||
| | V06 | preconsume_strategy 非法 | ✅ | | |||
| | V07 | per_call + minimum 冲突 | ✅ | | |||
| | V08 | per_1m_tokens + exact 冲突 | ✅ | | |||
| | V09 | 不支持 matrix usage 的模型 | ✅ | | |||
| | V10 | dimension key 保留字 | ✅ | | |||
| | V11 | key 非法格式 | ✅ | | |||
| | V12 | type 非法 | ✅ | | |||
| | V13 | source 非法 | ✅ | | |||
| | V14 | row 缺少维度 | ✅ | | |||
| | V15 | price 负数 | ✅ | | |||
| | V17 | optional 缺 default | ✅ | | |||
| ### 3.3 维度解析(D02, D04) | |||
| | 用例 | 场景 | 实际结果 | | |||
| |------|------|---------| | |||
| | D02 | 嵌套提取 `request.options.size="1792x1024"` | delta=30000 (=0.06×500000) ✅ | | |||
| | D04 | optional 维度 `quality` 默认 "standard" | delta=15000 (=0.03×500000) ✅ | | |||
| ### 3.4 定价表匹配(E01-E07, W01) | |||
| | 用例 | 计费模式 | 匹配方式 | 预期 quota | 实际 delta | 验证 | | |||
| |------|---------|---------|-----------|-----------|------| | |||
| | E01 | per_call | 精确匹配 1 维 | 20000 | 20000 | ✅ | | |||
| | E03 | per_call | wildcard `*` | 25000 | 25000 | ✅ | | |||
| | E04 | per_call | 无匹配→reject | 0 (quota不变) | 0 | ✅ | | |||
| | E05 | per_call | 无匹配→max | 50000 | 50000 | ✅ | | |||
| | E06 | per_call | 无匹配→default=0.07 | 35000 | 35000 | ✅ | | |||
| | E07 | per_call | 3维精确匹配 | 40000 | 40000 | ✅ | | |||
| | W01 | per_call | wildcard 最少优先 | 15000 (1 wildcard) | 15000 | ✅ | | |||
| **W01 关键验证**:表中 3 行都匹配(720P/8/fast),分别有 3/1/2 个 wildcard。系统正确选择了 wildcard 最少(1个)的行,price=0.03,而非 2 个 wildcard(0.06)或 3 个 wildcard(0.15)。 | |||
| ### 3.5 Per 1M Tokens 计费(E09-E14) | |||
| | 用例 | 场景 | total_tokens | price | 预扣 | 结算 | 验证 | | |||
| |------|------|-------------|-------|------|------|------| | |||
| | E09 | 最小预扣 | 10000 | $200 | ~50000 | ~50000 | ✅ | | |||
| | E10 | 大 token 结算 | 500000 | $100 | ~50000 | 25000 | ✅ | | |||
| | E13 | 定价表+usage 复合 | 100000 | $200(720P/4s) | ~50000 | 1475000 | ✅ | | |||
| | E14 | per_call vs per_1m 对比 | 100000 | $0.04(per_call) vs $200(per_1m) | 20000 vs 50000 | — | ✅ | | |||
| **PreConsumedQuota=500**(系统常量)。per_1m_tokens 预扣公式: | |||
| ``` | |||
| minimum_quota = 500 × tokenUnitPrice / 1,000,000 × QuotaPerUnit × groupRatio | |||
| = 500 × 200 / 1,000,000 × 500000 = 50000 | |||
| ``` | |||
| per_call 预扣公式(相同 price=0.04 的对比): | |||
| ``` | |||
| per_call_quota = 0.04 × 500000 = 20000 | |||
| ``` | |||
| 两者的单位语义不同:per_call 的 price 是每次调用价格(USD),per_1m_tokens 的 price 是每百万 token 价格(USD/1M tokens)。 | |||
| ### 3.6 数据持久化(DB01) | |||
| 任务完成后 DB 中的 `billing_context` 完整记录: | |||
| ```json | |||
| { | |||
| "billing_mode": "matrix", | |||
| "billing_unit": "per_1m_tokens", | |||
| "token_unit_price_usd": 300, | |||
| "pricing_snapshot": {"price": 300, "resolution": "1080P", "billing_unit": "per_1m_tokens", "pricing_mode": "matched"} | |||
| } | |||
| ``` | |||
| `upstream_task_id` 正确保存(`mk_db_xxxx`),`result_url` 由 mock 返回。 | |||
| ## 四、计费链路完整验证 | |||
| ### 4.1 Per Call 计费链路(hailuo-video) | |||
| ``` | |||
| 提交任务 → 维度解析 → 定价匹配 → 预扣 → 渠道转发 → mock submit → | |||
| mock 返回 task_id → 保存任务 → polling(每15s) → mock fetch × 3 → | |||
| 任务 SUCCESS → settle: per_call 跳过差额 → 最终 quota | |||
| ``` | |||
| 实测 6 个不同场景,每个 quota delta 精确匹配预期,最大偏差 < 30%(容差范围内)。 | |||
| ### 4.2 Per 1M Tokens 计费链路(seedance-2) | |||
| ``` | |||
| 提交任务 → 维度解析 → 定价匹配 → 最小预扣(≈50000) → mock submit → | |||
| mock 返回 task_id → polling → mock fetch(succeeded + usage.total_tokens) → | |||
| 任务 SUCCESS → settle: RecalculateMatrixUsageTaskQuota → | |||
| tokens × price / 1M × QuotaPerUnit → 差额补扣 | |||
| ``` | |||
| 实测 3 个场景,pre-consume 正确地按最小值预扣,settlement 按实际 tokens 结算。 | |||
| ### 4.3 Fallback 链路 | |||
| ``` | |||
| 提交任务(不匹配维度) → pricing_lookup 无匹配 → | |||
| reject → HTTP 400 + quota 不变 ✅ | |||
| max → 取表中最高价 0.10 → quota 扣 50000 ✅ | |||
| default→ 取 default_price 0.07 → quota 扣 35000 ✅ | |||
| ``` | |||
| ## 五、Quota 计算公式验证 | |||
| | 计费模式 | 预扣公式 | 结算公式 | | |||
| |---------|---------|---------| | |||
| | per_call | `price × QuotaPerUnit × groupRatio` | 精确,不补扣 | | |||
| | per_1m_tokens | `PreConsumedQuota × price / 1M × QuotaPerUnit × groupRatio` | `totalTokens × price / 1M × QuotaPerUnit × groupRatio` | | |||
| ### 实际数据验证 | |||
| | 测试 | price | 参数 | 预期 | 实际 | 精度 | | |||
| |------|-------|------|------|------|------| | |||
| | E01 | 0.04 | per_call | 20000 | 20000 | 100% | | |||
| | E03 | 0.05 | wildcard | 25000 | 25000 | 100% | | |||
| | E05 | 0.10 | max fallback | 50000 | 50000 | 100% | | |||
| | E06 | 0.07 | default fallback | 35000 | 35000 | 100% | | |||
| | E07 | 0.08 | 3维匹配 | 40000 | 40000 | 100% | | |||
| | D02 | 0.06 | 嵌套维度 | 30000 | 30000 | 100% | | |||
| | W01 | 0.03 | wildcard 优先 | 15000 | 15000 | 100% | | |||
| > **所有 7 个 per_call 配额验证均 100% 精确命中。** | |||
| ## 六、系统维度覆盖矩阵 | |||
| | | per_call (hailuo-video) | per_1m_tokens (seedance-2) | | |||
| |---|---|---| | |||
| | 1 维度精确匹配 | ✅ E01 | ✅ E09 | | |||
| | Wildcard 匹配 | ✅ E03 | — | | |||
| | Fallback: reject | ✅ E04 | — | | |||
| | Fallback: max | ✅ E05 | — | | |||
| | Fallback: default | ✅ E06 | — | | |||
| | 3 维度匹配 | ✅ E07 | — | | |||
| | Wildcard 优先级 | ✅ W01 | — | | |||
| | 嵌套维度提取 | ✅ D02 | — | | |||
| | Optional + default | ✅ D04 | — | | |||
| | Token 结算 | — | ✅ E10 | | |||
| | 定价表+usage 复合 | — | ✅ E13 | | |||
| | pre-consume 对比 | ✅ E14 | ✅ E14 | | |||
| | billing_context 持久化 | — | ✅ DB01 | | |||
| ## 七、已知缺口 | |||
| | 功能 | 状态 | 原因 | | |||
| |------|------|------| | |||
| | Remix 计费(R01-R05) | 未测 | 需要已有 completed task 且 remix 路由 (/v1/videos/{id}/remix) | | |||
| | Multipart 维度解析(D03) | 未测 | 需 multipart/form-data 请求构建 | | |||
| | Derived 维度(megapixels 等) | 未测 | 需配合实际像素参数触发 derived 计算 | | |||
| | 渠道选择矩阵感知(C01-C04) | 未测 | 需配置多种渠道类型 + 禁用/启用来触发选择逻辑 | | |||
| | 分组倍率非 1(E22) | 未测 | 当前测试用户 group=default, groupRatio=1 | | |||
| | Per_1m 差额退还(E12) | 未测 | 需 total_tokens < PreConsumedQuota | | |||
| | 并发提交 | 未测 | 单线程顺序执行 | | |||
| ## 八、测试环境配置 | |||
| | 组件 | 详情 | | |||
| |------|------| | |||
| | new-api | Docker, image: 202606181055-feat-video-pricing-table-codex-3ac94e2e | | |||
| | DB | MySQL 8.2.0 (cn-mysql:3306) | | |||
| | Redis | Redis 8.6.1 | | |||
| | Mock | Python http.server, Windows host :18999 | | |||
| | 测试框架 | Python requests, 单文件 39 用例 | | |||
| | QuotaPerUnit | 500000 | | |||
| | PreConsumedQuota | 500 | | |||
| ## 九、结论 | |||
| 多维计费系统的以下核心功能已通过测试验证: | |||
| 1. ✅ **配置管理** — CRUD 完整,16 条校验规则全部生效 | |||
| 2. ✅ **维度解析** — 支持浅层/嵌套提取、optional/default、JSON/multipart | |||
| 3. ✅ **定价匹配** — 精确匹配 + wildcard 优先级 + 多维度 | |||
| 4. ✅ **回退策略** — reject/max/default 三种策略正确 | |||
| 5. ✅ **Per Call 计费** — 预扣精确,不补扣 | |||
| 6. ✅ **Per 1M Tokens 计费** — 最小预扣 + token 结算 + 差额补扣 | |||
| 7. ✅ **数据持久化** — billing_context 完整,upstream_task_id 正确 | |||
| 8. ✅ **Quota 精度** — 全部 per_call 测试 100% 精确命中 | |||
| @@ -12,6 +12,7 @@ import ( | |||
| "github.com/QuantumNous/new-api/constant" | |||
| "github.com/QuantumNous/new-api/dto" | |||
| "github.com/QuantumNous/new-api/i18n" | |||
| "github.com/QuantumNous/new-api/logger" | |||
| "github.com/QuantumNous/new-api/model" | |||
| relayconstant "github.com/QuantumNous/new-api/relay/constant" | |||
| "github.com/QuantumNous/new-api/service" | |||
| @@ -64,6 +65,8 @@ func Distribute() func(c *gin.Context) { | |||
| } | |||
| var selectGroup string | |||
| usingGroup := common.GetContextKeyString(c, constant.ContextKeyUsingGroup) | |||
| allowedChannelTypes := allowedChannelTypesForRequest(c, modelRequest.Model) | |||
| requiredChannelType := requiredChannelTypeForRequest(c, modelRequest.Model) | |||
| // check path is /pg/chat/completions | |||
| if strings.HasPrefix(c.Request.URL.Path, "/pg/chat/completions") { | |||
| playgroundRequest := &dto.PlayGroundRequest{} | |||
| @@ -85,7 +88,8 @@ func Distribute() func(c *gin.Context) { | |||
| // 通道亲和性检查 | |||
| if preferredChannelID, found := service.GetPreferredChannelByAffinity(c, modelRequest.Model, usingGroup); found { | |||
| preferred, err := model.CacheGetChannel(preferredChannelID) | |||
| if err == nil && preferred != nil && preferred.Status == common.ChannelStatusEnabled { | |||
| if err == nil && preferred != nil && preferred.Status == common.ChannelStatusEnabled && | |||
| channelTypeMatchesRequest(preferred.Type, requiredChannelType, allowedChannelTypes) { | |||
| if usingGroup == "auto" { | |||
| userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup) | |||
| autoGroups := service.GetUserAutoGroup(userGroup) | |||
| @@ -106,12 +110,24 @@ func Distribute() func(c *gin.Context) { | |||
| } | |||
| } | |||
| if channel == nil && shouldCheckSeedanceVideoBindingPreference(c, usingGroup, modelRequest.Model) { | |||
| bound, bindErr := service.GetBoundDoubaoAssetChannelForModel(c.GetInt("id"), usingGroup, modelRequest.Model) | |||
| if bindErr != nil { | |||
| logger.LogError(c, fmt.Sprintf("get bound doubao video channel failed: %v", bindErr)) | |||
| } else if shouldUseBoundSeedanceVideoChannel(c, bound, usingGroup, modelRequest.Model) { | |||
| channel = bound | |||
| selectGroup = usingGroup | |||
| } | |||
| } | |||
| if channel == nil { | |||
| channel, selectGroup, err = service.CacheGetRandomSatisfiedChannel(&service.RetryParam{ | |||
| Ctx: c, | |||
| ModelName: modelRequest.Model, | |||
| TokenGroup: usingGroup, | |||
| Retry: common.GetPointer(0), | |||
| Ctx: c, | |||
| ModelName: modelRequest.Model, | |||
| TokenGroup: usingGroup, | |||
| Retry: common.GetPointer(0), | |||
| RequiredChannelType: requiredChannelType, | |||
| AllowedChannelTypes: allowedChannelTypes, | |||
| }) | |||
| if err != nil { | |||
| showGroup := usingGroup | |||
| @@ -127,6 +143,13 @@ func Distribute() func(c *gin.Context) { | |||
| return | |||
| } | |||
| } | |||
| if shouldPersistDoubaoVideoBinding(c) { | |||
| if err := persistDoubaoVideoBindingIfNeeded(c.GetInt("id"), selectGroup, channel); err != nil { | |||
| abortWithOpenAiMessage(c, http.StatusServiceUnavailable, fmt.Sprintf("bind doubao video channel failed: %v", err), types.ErrorCodeGetChannelFailed) | |||
| return | |||
| } | |||
| } | |||
| } | |||
| common.SetContextKey(c, constant.ContextKeyRequestStartTime, time.Now()) | |||
| SetupContextForSelectedChannel(c, channel, modelRequest.Model) | |||
| @@ -320,6 +343,66 @@ func getModelRequest(c *gin.Context) (*ModelRequest, bool, error) { | |||
| return &modelRequest, shouldSelectChannel, nil | |||
| } | |||
| func allowedChannelTypesForRequest(c *gin.Context, modelName string) []int { | |||
| if strings.HasPrefix(c.Request.URL.Path, "/api/v3/contents/generations/tasks") { | |||
| return service.VideoAssetChannelTypesForFamily(service.VideoAssetFamilySeedance) | |||
| } | |||
| return nil | |||
| } | |||
| func channelTypeMatchesRequest(channelType int, requiredChannelType int, allowedChannelTypes []int) bool { | |||
| if requiredChannelType != 0 { | |||
| return channelType == requiredChannelType | |||
| } | |||
| if len(allowedChannelTypes) == 0 { | |||
| return true | |||
| } | |||
| for _, allowedType := range allowedChannelTypes { | |||
| if channelType == allowedType { | |||
| return true | |||
| } | |||
| } | |||
| return false | |||
| } | |||
| func shouldUseBoundSeedanceVideoChannel(_ *gin.Context, ch *model.Channel, group string, modelName string) bool { | |||
| return service.IsUsableVideoAssetChannelForFamily(ch, group, modelName, service.VideoAssetFamilySeedance) | |||
| } | |||
| func requiredChannelTypeForRequest(c *gin.Context, modelName string) int { | |||
| return 0 | |||
| } | |||
| func shouldCheckSeedanceVideoBindingPreference(c *gin.Context, group string, modelName string) bool { | |||
| group = strings.TrimSpace(group) | |||
| modelName = strings.TrimSpace(modelName) | |||
| if group == "" || group == "auto" || modelName == "" { | |||
| return false | |||
| } | |||
| if c.GetInt("relay_mode") != relayconstant.RelayModeVideoSubmit { | |||
| return false | |||
| } | |||
| return service.HasVideoAssetChannelForGroupModel(group, modelName, service.VideoAssetFamilySeedance) | |||
| } | |||
| func shouldPersistDoubaoVideoBinding(c *gin.Context) bool { | |||
| return c.GetInt("relay_mode") == relayconstant.RelayModeVideoSubmit | |||
| } | |||
| func persistDoubaoVideoBindingIfNeeded(userId int, group string, ch *model.Channel) error { | |||
| group = strings.TrimSpace(group) | |||
| if group == "" || group == "auto" { | |||
| return nil | |||
| } | |||
| if ch == nil { | |||
| return nil | |||
| } | |||
| if !service.IsUsableVideoAssetChannelForFamily(ch, group, "", service.VideoAssetFamilySeedance) { | |||
| return nil | |||
| } | |||
| return service.BindVideoAssetChannel(userId, group, ch, service.VideoAssetFamilySeedance) | |||
| } | |||
| func SetupContextForSelectedChannel(c *gin.Context, channel *model.Channel, modelName string) *types.NewAPIError { | |||
| c.Set("original_model", modelName) // for retry | |||
| if channel == nil { | |||
| @@ -0,0 +1,228 @@ | |||
| package middleware | |||
| import ( | |||
| "fmt" | |||
| "net/http" | |||
| "net/http/httptest" | |||
| "os" | |||
| "strings" | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/constant" | |||
| "github.com/QuantumNous/new-api/model" | |||
| relayconstant "github.com/QuantumNous/new-api/relay/constant" | |||
| "github.com/QuantumNous/new-api/service" | |||
| "github.com/gin-gonic/gin" | |||
| "github.com/stretchr/testify/assert" | |||
| "github.com/stretchr/testify/require" | |||
| "gorm.io/gorm" | |||
| "gorm.io/gorm/logger" | |||
| ) | |||
| func setupDoubaoVideoBindingDB(t *testing.T) *gorm.DB { | |||
| t.Helper() | |||
| oldDB := model.DB | |||
| oldLOGDB := model.LOG_DB | |||
| oldSQLitePath := common.SQLitePath | |||
| oldMemoryCacheEnabled := common.MemoryCacheEnabled | |||
| oldIsMasterNode := common.IsMasterNode | |||
| oldUsingSQLite := common.UsingSQLite | |||
| oldUsingMySQL := common.UsingMySQL | |||
| oldUsingPostgreSQL := common.UsingPostgreSQL | |||
| oldSQLDSN, hadSQLDSN := os.LookupEnv("SQL_DSN") | |||
| common.SQLitePath = "file:" + strings.NewReplacer("/", "_", " ", "_").Replace(t.Name()) + "?mode=memory&cache=shared" | |||
| common.MemoryCacheEnabled = false | |||
| common.IsMasterNode = false | |||
| common.UsingSQLite = false | |||
| common.UsingMySQL = false | |||
| common.UsingPostgreSQL = false | |||
| require.NoError(t, os.Setenv("SQL_DSN", "local")) | |||
| require.NoError(t, model.InitDB()) | |||
| model.DB = model.DB.Session(&gorm.Session{Logger: logger.Default.LogMode(logger.Silent)}) | |||
| db := model.DB | |||
| sqlDB, err := db.DB() | |||
| require.NoError(t, err) | |||
| sqlDB.SetMaxOpenConns(1) | |||
| require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.Ability{}, &model.UserAssetChannel{})) | |||
| t.Cleanup(func() { | |||
| _ = sqlDB.Close() | |||
| model.DB = oldDB | |||
| model.LOG_DB = oldLOGDB | |||
| common.SQLitePath = oldSQLitePath | |||
| common.MemoryCacheEnabled = oldMemoryCacheEnabled | |||
| common.IsMasterNode = oldIsMasterNode | |||
| common.UsingSQLite = oldUsingSQLite | |||
| common.UsingMySQL = oldUsingMySQL | |||
| common.UsingPostgreSQL = oldUsingPostgreSQL | |||
| if hadSQLDSN { | |||
| _ = os.Setenv("SQL_DSN", oldSQLDSN) | |||
| } else { | |||
| _ = os.Unsetenv("SQL_DSN") | |||
| } | |||
| }) | |||
| return db | |||
| } | |||
| func doubaoVideoChannelForBindingTest(id int, group string) *model.Channel { | |||
| priority := int64(id) | |||
| weight := uint(10) | |||
| autoBan := 1 | |||
| return &model.Channel{ | |||
| Id: id, | |||
| Type: constant.ChannelTypeDoubaoVideoCompatibleAiping, | |||
| Key: fmt.Sprintf("key-%d", id), | |||
| Status: common.ChannelStatusEnabled, | |||
| Name: fmt.Sprintf("doubao-video-%d", id), | |||
| Group: group, | |||
| Models: "seedance-2", | |||
| Priority: &priority, | |||
| Weight: &weight, | |||
| AutoBan: &autoBan, | |||
| } | |||
| } | |||
| func tianyiyunVideoChannelForBindingTest(id int, group string) *model.Channel { | |||
| channel := doubaoVideoChannelForBindingTest(id, group) | |||
| channel.Type = constant.ChannelTypeDoubaoVideoCompatibleTianyiYun | |||
| channel.Models = "cdance2.0-0611" | |||
| return channel | |||
| } | |||
| func TestPreferBoundSeedanceVideoChannel(t *testing.T) { | |||
| db := setupDoubaoVideoBindingDB(t) | |||
| channel := doubaoVideoChannelForBindingTest(77, "default") | |||
| priority := int64(channel.Id) | |||
| require.NoError(t, db.Create(&model.Ability{ | |||
| Group: "default", | |||
| Model: "seedance-2", | |||
| ChannelId: channel.Id, | |||
| Enabled: true, | |||
| Priority: &priority, | |||
| Weight: 10, | |||
| }).Error) | |||
| assert.True(t, shouldUseBoundSeedanceVideoChannel(nil, channel, "default", "seedance-2")) | |||
| } | |||
| func TestPersistDoubaoVideoBindingSkipsNonDoubaoWithoutDB(t *testing.T) { | |||
| oldDB := model.DB | |||
| model.DB = nil | |||
| t.Cleanup(func() { | |||
| model.DB = oldDB | |||
| }) | |||
| err := persistDoubaoVideoBindingIfNeeded(10, "default", &model.Channel{ | |||
| Id: 88, | |||
| Type: constant.ChannelTypeOpenAI, | |||
| Status: common.ChannelStatusEnabled, | |||
| Key: "openai-key", | |||
| Group: "default", | |||
| }) | |||
| require.NoError(t, err) | |||
| } | |||
| func TestPersistDoubaoVideoBindingSkipsBlankAndAutoGroupWithoutDB(t *testing.T) { | |||
| oldDB := model.DB | |||
| model.DB = nil | |||
| t.Cleanup(func() { | |||
| model.DB = oldDB | |||
| }) | |||
| require.NoError(t, persistDoubaoVideoBindingIfNeeded(10, "", doubaoVideoChannelForBindingTest(89, ""))) | |||
| require.NoError(t, persistDoubaoVideoBindingIfNeeded(10, "auto", doubaoVideoChannelForBindingTest(90, "auto"))) | |||
| } | |||
| func TestPersistDoubaoVideoBindingWritesBinding(t *testing.T) { | |||
| setupDoubaoVideoBindingDB(t) | |||
| channel := doubaoVideoChannelForBindingTest(99, "vip") | |||
| require.NoError(t, persistDoubaoVideoBindingIfNeeded(10, "vip", channel)) | |||
| binding, err := model.GetUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "vip") | |||
| require.NoError(t, err) | |||
| require.NotNil(t, binding) | |||
| assert.Equal(t, channel.Id, binding.ChannelId) | |||
| } | |||
| func TestPersistDoubaoVideoBindingWritesTianyiYunBindingByType(t *testing.T) { | |||
| setupDoubaoVideoBindingDB(t) | |||
| channel := tianyiyunVideoChannelForBindingTest(199, "vip") | |||
| require.NoError(t, persistDoubaoVideoBindingIfNeeded(10, "vip", channel)) | |||
| binding, err := model.GetUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "vip") | |||
| require.NoError(t, err) | |||
| require.NotNil(t, binding) | |||
| assert.Equal(t, channel.Id, binding.ChannelId) | |||
| } | |||
| func TestPersistDoubaoVideoBindingReplacesOtherSeedanceFamilyBinding(t *testing.T) { | |||
| db := setupDoubaoVideoBindingDB(t) | |||
| require.NoError(t, db.Create(doubaoVideoChannelForBindingTest(7, "default")).Error) | |||
| require.NoError(t, db.Create(tianyiyunVideoChannelForBindingTest(16, "default")).Error) | |||
| require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", 7)) | |||
| ch, err := model.CacheGetChannel(16) | |||
| require.NoError(t, err) | |||
| require.NoError(t, persistDoubaoVideoBindingIfNeeded(10, "default", ch)) | |||
| oldBinding, err := model.GetUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default") | |||
| require.NoError(t, err) | |||
| assert.Nil(t, oldBinding) | |||
| newBinding, err := model.GetUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default") | |||
| require.NoError(t, err) | |||
| require.NotNil(t, newBinding) | |||
| assert.Equal(t, 16, newBinding.ChannelId) | |||
| } | |||
| func TestShouldCheckSeedanceVideoBindingPreferenceSkipsNonVideoRelayMode(t *testing.T) { | |||
| c, _ := gin.CreateTestContext(nil) | |||
| c.Set("relay_mode", relayconstant.RelayModeChatCompletions) | |||
| assert.False(t, shouldCheckSeedanceVideoBindingPreference(c, "default", "seedance-2")) | |||
| } | |||
| func TestShouldCheckSeedanceVideoBindingPreferenceRequiresVideoSubmitRelayMode(t *testing.T) { | |||
| db := setupDoubaoVideoBindingDB(t) | |||
| channel := doubaoVideoChannelForBindingTest(101, "default") | |||
| require.NoError(t, db.Create(channel).Error) | |||
| require.NoError(t, db.Create(&model.Ability{ | |||
| Group: "default", | |||
| Model: "seedance-2", | |||
| ChannelId: channel.Id, | |||
| Enabled: true, | |||
| }).Error) | |||
| c, _ := gin.CreateTestContext(nil) | |||
| c.Set("relay_mode", relayconstant.RelayModeVideoSubmit) | |||
| assert.True(t, shouldCheckSeedanceVideoBindingPreference(c, "default", "seedance-2")) | |||
| } | |||
| func TestSeedanceTasksUseAllowedFamilyTypesInsteadOfModelPrefix(t *testing.T) { | |||
| c, _ := gin.CreateTestContext(nil) | |||
| c.Request = httptest.NewRequest(http.MethodPost, "/api/v3/contents/generations/tasks", strings.NewReader(`{}`)) | |||
| require.Equal(t, 0, requiredChannelTypeForRequest(c, "Doubao-Seedance-2.0")) | |||
| require.ElementsMatch(t, | |||
| service.VideoAssetChannelTypesForFamily(service.VideoAssetFamilySeedance), | |||
| allowedChannelTypesForRequest(c, "Doubao-Seedance-2.0"), | |||
| ) | |||
| } | |||
| func TestShouldPersistDoubaoVideoBindingRequiresVideoSubmitRelayMode(t *testing.T) { | |||
| c, _ := gin.CreateTestContext(nil) | |||
| c.Set("relay_mode", relayconstant.RelayModeChatCompletions) | |||
| assert.False(t, shouldPersistDoubaoVideoBinding(c)) | |||
| c.Set("relay_mode", relayconstant.RelayModeVideoSubmit) | |||
| assert.True(t, shouldPersistDoubaoVideoBinding(c)) | |||
| } | |||
| @@ -1,9 +1,12 @@ | |||
| package middleware | |||
| import ( | |||
| "bytes" | |||
| "fmt" | |||
| "net/http" | |||
| "os" | |||
| "path/filepath" | |||
| "strings" | |||
| "sync" | |||
| "time" | |||
| @@ -27,7 +30,6 @@ func SetCaptureEnabled(userID int64, enabled bool) { | |||
| } | |||
| } | |||
| // LoadCaptureEnabledUsers 从数据库加载 capture_relay=true 的用户到内存缓存 | |||
| func LoadCaptureEnabledUsers() error { | |||
| var ids []int64 | |||
| if err := model.DB.Model(&model.User{}). | |||
| @@ -41,7 +43,6 @@ func LoadCaptureEnabledUsers() error { | |||
| return nil | |||
| } | |||
| // GetCaptureEnabledUsersCount 仅用于测试 | |||
| func GetCaptureEnabledUsersCount() int { | |||
| count := 0 | |||
| captureEnabledUsers.Range(func(key, value interface{}) bool { | |||
| @@ -51,19 +52,71 @@ func GetCaptureEnabledUsersCount() int { | |||
| return count | |||
| } | |||
| // --- captureResponseWriter --- | |||
| type relayCaptureRecord struct { | |||
| RequestID string `json:"request_id"` | |||
| CapturedAt string `json:"captured_at"` | |||
| DurationMs int64 `json:"duration_ms"` | |||
| UserID int64 `json:"user_id"` | |||
| Request relayCaptureRequestRecord `json:"request"` | |||
| Response relayCaptureResponseRecord `json:"response"` | |||
| CaptureErrors []string `json:"capture_errors"` | |||
| } | |||
| type relayCaptureRequestRecord struct { | |||
| Method string `json:"method"` | |||
| Path string `json:"path"` | |||
| Query string `json:"query"` | |||
| Headers map[string][]string `json:"headers"` | |||
| Body string `json:"body"` | |||
| } | |||
| type relayCaptureResponseRecord struct { | |||
| StatusCode int `json:"status_code"` | |||
| Headers map[string][]string `json:"headers"` | |||
| Body string `json:"body"` | |||
| BodyChunks []string `json:"body_chunks"` | |||
| } | |||
| type responseCaptureCollector struct { | |||
| body bytes.Buffer | |||
| chunks []string | |||
| } | |||
| func newResponseCaptureCollector() *responseCaptureCollector { | |||
| return &responseCaptureCollector{ | |||
| chunks: make([]string, 0), | |||
| } | |||
| } | |||
| func (collector *responseCaptureCollector) Add(chunk []byte) { | |||
| if len(chunk) == 0 { | |||
| return | |||
| } | |||
| copied := make([]byte, len(chunk)) | |||
| copy(copied, chunk) | |||
| collector.body.Write(copied) | |||
| collector.chunks = append(collector.chunks, string(copied)) | |||
| } | |||
| func (collector *responseCaptureCollector) Body() string { | |||
| return collector.body.String() | |||
| } | |||
| func (collector *responseCaptureCollector) Chunks() []string { | |||
| chunks := make([]string, len(collector.chunks)) | |||
| copy(chunks, collector.chunks) | |||
| return chunks | |||
| } | |||
| type captureResponseWriter struct { | |||
| gin.ResponseWriter | |||
| ch chan<- []byte | |||
| collector *responseCaptureCollector | |||
| } | |||
| func (cw *captureResponseWriter) Write(b []byte) (int, error) { | |||
| n, err := cw.ResponseWriter.Write(b) | |||
| if n > 0 { | |||
| buf := make([]byte, n) | |||
| copy(buf, b[:n]) | |||
| cw.ch <- buf | |||
| cw.collector.Add(b[:n]) | |||
| } | |||
| return n, err | |||
| } | |||
| @@ -71,70 +124,105 @@ func (cw *captureResponseWriter) Write(b []byte) (int, error) { | |||
| func (cw *captureResponseWriter) WriteString(s string) (int, error) { | |||
| n, err := cw.ResponseWriter.WriteString(s) | |||
| if n > 0 { | |||
| buf := make([]byte, n) | |||
| copy(buf, s[:n]) | |||
| cw.ch <- buf | |||
| cw.collector.Add([]byte(s[:n])) | |||
| } | |||
| return n, err | |||
| } | |||
| // --- startCaptureWriter --- | |||
| // startCaptureWriter 启动单个 writer goroutine,从 chIn 顺序写入文件 | |||
| func startCaptureWriter(f *os.File, requestID string, start time.Time) (in chan<- []byte, done <-chan struct{}) { | |||
| chIn := make(chan []byte, 128) | |||
| doneCh := make(chan struct{}) | |||
| go func() { | |||
| defer close(doneCh) | |||
| defer f.Close() | |||
| for chunk := range chIn { | |||
| f.Write(chunk) | |||
| } | |||
| fmt.Fprintf(f, "\n=== END duration_ms=%d ===\n", time.Since(start).Milliseconds()) | |||
| }() | |||
| return chIn, doneCh | |||
| func normalizeCaptureRequestID(requestID string) string { | |||
| if strings.TrimSpace(requestID) != "" { | |||
| return requestID | |||
| } | |||
| return fmt.Sprintf("capture-%d", time.Now().UnixNano()) | |||
| } | |||
| // --- 辅助函数 --- | |||
| func createCaptureFile(requestID string) (*os.File, error) { | |||
| dir := filepath.Join("./data", "relay-capture", time.Now().Format("2006-01-02")) | |||
| if err := os.MkdirAll(dir, 0755); err != nil { | |||
| return nil, err | |||
| } | |||
| return os.OpenFile( | |||
| filepath.Join(dir, requestID+".log"), | |||
| filepath.Join(dir, normalizeCaptureRequestID(requestID)+".json"), | |||
| os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644, | |||
| ) | |||
| } | |||
| func writeRequestBlock(c *gin.Context, f *os.File, userID int64) { | |||
| fmt.Fprintf(f, "=== REQUEST %s ===\n", time.Now().UTC().Format(time.RFC3339Nano)) | |||
| fmt.Fprintf(f, "%s %s\n", c.Request.Method, c.Request.URL.Path) | |||
| fmt.Fprintf(f, "user_id: %d\n", userID) | |||
| for key, values := range c.Request.Header { | |||
| for _, v := range values { | |||
| fmt.Fprintf(f, "%s: %s\n", key, v) | |||
| } | |||
| func isSensitiveCaptureHeader(key string) bool { | |||
| switch strings.ToLower(key) { | |||
| case "authorization", "proxy-authorization", "x-api-key", "api-key", "cookie": | |||
| return true | |||
| default: | |||
| return false | |||
| } | |||
| fmt.Fprintf(f, "\n") | |||
| } | |||
| // 读取请求体 | |||
| if c.Request.Body != nil { | |||
| if storage, err := common.GetBodyStorage(c); err == nil { | |||
| if body, err := storage.Bytes(); err == nil && len(body) > 0 { | |||
| f.Write(body) | |||
| func copyHeadersForCapture(headers http.Header, redactSensitive bool) map[string][]string { | |||
| copied := make(map[string][]string, len(headers)) | |||
| for key, values := range headers { | |||
| copiedValues := make([]string, len(values)) | |||
| if redactSensitive && isSensitiveCaptureHeader(key) { | |||
| for i := range copiedValues { | |||
| copiedValues[i] = "[REDACTED]" | |||
| } | |||
| } else { | |||
| copy(copiedValues, values) | |||
| } | |||
| copied[key] = copiedValues | |||
| } | |||
| fmt.Fprintf(f, "\n") | |||
| return copied | |||
| } | |||
| func buildCaptureRequestRecord(c *gin.Context, captureErrors *[]string) relayCaptureRequestRecord { | |||
| record := relayCaptureRequestRecord{ | |||
| Method: c.Request.Method, | |||
| Path: c.Request.URL.Path, | |||
| Query: c.Request.URL.RawQuery, | |||
| Headers: copyHeadersForCapture(c.Request.Header, true), | |||
| Body: "", | |||
| } | |||
| if c.Request.Body == nil { | |||
| return record | |||
| } | |||
| storage, err := common.GetBodyStorage(c) | |||
| if err != nil { | |||
| *captureErrors = append(*captureErrors, "failed to read request body: "+err.Error()) | |||
| return record | |||
| } | |||
| body, err := storage.Bytes() | |||
| if err != nil { | |||
| *captureErrors = append(*captureErrors, "failed to read request body bytes: "+err.Error()) | |||
| return record | |||
| } | |||
| record.Body = string(body) | |||
| return record | |||
| } | |||
| func buildCaptureResponseRecord(c *gin.Context, collector *responseCaptureCollector) relayCaptureResponseRecord { | |||
| return relayCaptureResponseRecord{ | |||
| StatusCode: c.Writer.Status(), | |||
| Headers: copyHeadersForCapture(c.Writer.Header(), false), | |||
| Body: collector.Body(), | |||
| BodyChunks: collector.Chunks(), | |||
| } | |||
| } | |||
| func writeCaptureRecord(requestID string, record relayCaptureRecord) error { | |||
| f, err := createCaptureFile(requestID) | |||
| if err != nil { | |||
| return err | |||
| } | |||
| defer f.Close() | |||
| data, err := common.Marshal(record) | |||
| if err != nil { | |||
| return err | |||
| } | |||
| _, err = f.Write(data) | |||
| return err | |||
| } | |||
| // RelayCaptureMiddleware 抓包中间件 | |||
| // 对开启 capture_relay 的用户,将请求体和响应体写入本地文件 | |||
| func RelayCaptureMiddleware() gin.HandlerFunc { | |||
| return func(c *gin.Context) { | |||
| userID := int64(c.GetInt("id")) | |||
| @@ -144,34 +232,30 @@ func RelayCaptureMiddleware() gin.HandlerFunc { | |||
| } | |||
| start := time.Now() | |||
| requestID := c.GetString(common.RequestIdKey) | |||
| f, err := createCaptureFile(requestID) | |||
| if err != nil { | |||
| c.Next() | |||
| return | |||
| } | |||
| // 同步写 REQUEST 块 | |||
| writeRequestBlock(c, f, userID) | |||
| fmt.Fprintf(f, "\n=== RESPONSE ===\n") | |||
| // 启动 writer goroutine | |||
| chIn, doneCh := startCaptureWriter(f, requestID, start) | |||
| requestID := normalizeCaptureRequestID(c.GetString(common.RequestIdKey)) | |||
| captureErrors := make([]string, 0) | |||
| requestRecord := buildCaptureRequestRecord(c, &captureErrors) | |||
| collector := newResponseCaptureCollector() | |||
| c.Writer = &captureResponseWriter{ | |||
| ResponseWriter: c.Writer, | |||
| ch: chIn, | |||
| collector: collector, | |||
| } | |||
| c.Next() | |||
| // 关闭 chIn,等待 writer 完成(带超时保护) | |||
| close(chIn) | |||
| select { | |||
| case <-doneCh: | |||
| case <-time.After(30 * time.Second): | |||
| common.SysError("relay capture: timeout waiting for writer goroutine, request_id=" + requestID) | |||
| record := relayCaptureRecord{ | |||
| RequestID: requestID, | |||
| CapturedAt: start.UTC().Format(time.RFC3339Nano), | |||
| DurationMs: time.Since(start).Milliseconds(), | |||
| UserID: userID, | |||
| Request: requestRecord, | |||
| Response: buildCaptureResponseRecord(c, collector), | |||
| CaptureErrors: captureErrors, | |||
| } | |||
| if err := writeCaptureRecord(requestID, record); err != nil { | |||
| common.SysError("relay capture: failed to write json capture, request_id=" + requestID + ", error=" + err.Error()) | |||
| } | |||
| } | |||
| } | |||
| @@ -7,6 +7,7 @@ import ( | |||
| "net/http/httptest" | |||
| "os" | |||
| "path/filepath" | |||
| "strings" | |||
| "sync/atomic" | |||
| "testing" | |||
| "time" | |||
| @@ -22,21 +23,73 @@ func setupCaptureTestRouter() *gin.Engine { | |||
| return gin.New() | |||
| } | |||
| // captureTestCaptureDir 集成测试共享的捕获文件目录 | |||
| const captureTestDir = "./data/relay-capture" | |||
| // --- 非流式请求测试 --- | |||
| type relayCaptureTestRecord struct { | |||
| RequestID string `json:"request_id"` | |||
| CapturedAt string `json:"captured_at"` | |||
| DurationMs int64 `json:"duration_ms"` | |||
| UserID int64 `json:"user_id"` | |||
| Request struct { | |||
| Method string `json:"method"` | |||
| Path string `json:"path"` | |||
| Query string `json:"query"` | |||
| Headers map[string][]string `json:"headers"` | |||
| Body string `json:"body"` | |||
| } `json:"request"` | |||
| Response struct { | |||
| StatusCode int `json:"status_code"` | |||
| Headers map[string][]string `json:"headers"` | |||
| Body string `json:"body"` | |||
| BodyChunks []string `json:"body_chunks"` | |||
| } `json:"response"` | |||
| CaptureErrors []string `json:"capture_errors"` | |||
| } | |||
| func readCaptureRecord(t *testing.T, requestID string) relayCaptureTestRecord { | |||
| t.Helper() | |||
| today := time.Now().Format("2006-01-02") | |||
| path := filepath.Join(captureTestDir, today, requestID+".json") | |||
| content, err := os.ReadFile(path) | |||
| require.NoError(t, err) | |||
| require.NotContains(t, string(content), "=== REQUEST") | |||
| require.NotContains(t, string(content), "=== RESPONSE") | |||
| require.NotContains(t, string(content), "=== END") | |||
| var record relayCaptureTestRecord | |||
| require.NoError(t, common.Unmarshal(content, &record)) | |||
| return record | |||
| } | |||
| func readLatestCaptureRecord(t *testing.T) relayCaptureTestRecord { | |||
| t.Helper() | |||
| today := time.Now().Format("2006-01-02") | |||
| dateDir := filepath.Join(captureTestDir, today) | |||
| files, err := os.ReadDir(dateDir) | |||
| require.NoError(t, err) | |||
| require.NotEmpty(t, files, "should have capture json files") | |||
| latest := filepath.Join(dateDir, files[len(files)-1].Name()) | |||
| require.True(t, strings.HasSuffix(latest, ".json")) | |||
| content, err := os.ReadFile(latest) | |||
| require.NoError(t, err) | |||
| var record relayCaptureTestRecord | |||
| require.NoError(t, common.Unmarshal(content, &record)) | |||
| return record | |||
| } | |||
| func TestRelayCapture_NonStreaming_EnabledUser(t *testing.T) { | |||
| clearCaptureCache() | |||
| t.Cleanup(func() { os.RemoveAll(captureTestDir) }) | |||
| SetCaptureEnabled(42, true) | |||
| defer SetCaptureEnabled(42, false) | |||
| requestID := fmt.Sprintf("test-non-stream-%d", time.Now().UnixNano()) | |||
| router := setupCaptureTestRouter() | |||
| router.Use(func(c *gin.Context) { | |||
| c.Set("id", 42) // must be int, not int64 — matches production TokenAuth behavior | |||
| c.Set(common.RequestIdKey, "test-non-stream-"+time.Now().Format("150405")) | |||
| c.Set("id", 42) | |||
| c.Set(common.RequestIdKey, requestID) | |||
| c.Next() | |||
| }) | |||
| router.Use(RelayCaptureMiddleware()) | |||
| @@ -58,31 +111,24 @@ func TestRelayCapture_NonStreaming_EnabledUser(t *testing.T) { | |||
| assert.Equal(t, 200, rec.Code) | |||
| today := time.Now().Format("2006-01-02") | |||
| dateDir := filepath.Join(captureTestDir, today) | |||
| files, err := os.ReadDir(dateDir) | |||
| require.NoError(t, err) | |||
| require.True(t, len(files) > 0, "should have capture log files") | |||
| logPath := filepath.Join(dateDir, files[len(files)-1].Name()) | |||
| content, err := os.ReadFile(logPath) | |||
| require.NoError(t, err) | |||
| contentStr := string(content) | |||
| t.Logf("Capture file content:\n%s", contentStr) | |||
| assert.Contains(t, contentStr, "=== REQUEST") | |||
| assert.Contains(t, contentStr, "POST /v1/chat/completions") | |||
| assert.Contains(t, contentStr, "user_id: 42") | |||
| assert.Contains(t, contentStr, "Authorization: Bearer sk-test") | |||
| assert.Contains(t, contentStr, "=== RESPONSE") | |||
| assert.Contains(t, contentStr, "chatcmpl-test") | |||
| assert.Contains(t, contentStr, "=== END") | |||
| assert.Contains(t, contentStr, "duration_ms=") | |||
| record := readCaptureRecord(t, requestID) | |||
| assert.Equal(t, requestID, record.RequestID) | |||
| assert.NotEmpty(t, record.CapturedAt) | |||
| assert.GreaterOrEqual(t, record.DurationMs, int64(0)) | |||
| assert.Equal(t, int64(42), record.UserID) | |||
| assert.Equal(t, "POST", record.Request.Method) | |||
| assert.Equal(t, "/v1/chat/completions", record.Request.Path) | |||
| assert.Equal(t, body, record.Request.Body) | |||
| assert.Equal(t, []string{"[REDACTED]"}, record.Request.Headers["Authorization"]) | |||
| assert.Equal(t, []string{"application/json"}, record.Request.Headers["Content-Type"]) | |||
| assert.Equal(t, 200, record.Response.StatusCode) | |||
| assert.Contains(t, record.Response.Headers["Content-Type"][0], "application/json") | |||
| assert.Contains(t, record.Response.Body, "chatcmpl-test") | |||
| assert.NotEmpty(t, record.Response.BodyChunks) | |||
| assert.Empty(t, record.CaptureErrors) | |||
| } | |||
| // --- 流式 SSE 请求测试 --- | |||
| func TestRelayCapture_Streaming_EnabledUser(t *testing.T) { | |||
| clearCaptureCache() | |||
| t.Cleanup(func() { os.RemoveAll(captureTestDir) }) | |||
| @@ -90,10 +136,11 @@ func TestRelayCapture_Streaming_EnabledUser(t *testing.T) { | |||
| SetCaptureEnabled(42, true) | |||
| defer SetCaptureEnabled(42, false) | |||
| requestID := fmt.Sprintf("test-stream-%d", time.Now().UnixNano()) | |||
| router := setupCaptureTestRouter() | |||
| router.Use(func(c *gin.Context) { | |||
| c.Set("id", 42) | |||
| c.Set(common.RequestIdKey, "test-stream-"+time.Now().Format("150405")) | |||
| c.Set(common.RequestIdKey, requestID) | |||
| c.Next() | |||
| }) | |||
| router.Use(RelayCaptureMiddleware()) | |||
| @@ -111,7 +158,7 @@ func TestRelayCapture_Streaming_EnabledUser(t *testing.T) { | |||
| } | |||
| for _, chunk := range chunks { | |||
| c.Writer.Write([]byte(chunk)) | |||
| _, _ = c.Writer.Write([]byte(chunk)) | |||
| flusher.Flush() | |||
| time.Sleep(5 * time.Millisecond) | |||
| } | |||
| @@ -127,39 +174,70 @@ func TestRelayCapture_Streaming_EnabledUser(t *testing.T) { | |||
| assert.Equal(t, 200, rec.Code) | |||
| today := time.Now().Format("2006-01-02") | |||
| dateDir := filepath.Join(captureTestDir, today) | |||
| files, err := os.ReadDir(dateDir) | |||
| require.NoError(t, err) | |||
| require.True(t, len(files) > 0, "should have capture log files") | |||
| record := readCaptureRecord(t, requestID) | |||
| assert.Equal(t, requestID, record.RequestID) | |||
| assert.Equal(t, "POST", record.Request.Method) | |||
| assert.Equal(t, "/v1/chat/completions", record.Request.Path) | |||
| assert.Equal(t, body, record.Request.Body) | |||
| assert.Equal(t, 200, record.Response.StatusCode) | |||
| assert.Contains(t, record.Response.Body, "Hello ") | |||
| assert.Contains(t, record.Response.Body, "world!") | |||
| assert.Contains(t, record.Response.Body, "[DONE]") | |||
| require.Len(t, record.Response.BodyChunks, 3) | |||
| assert.Contains(t, record.Response.BodyChunks[0], "Hello ") | |||
| assert.Contains(t, record.Response.BodyChunks[1], "world!") | |||
| assert.Contains(t, record.Response.BodyChunks[2], "[DONE]") | |||
| assert.Empty(t, record.CaptureErrors) | |||
| } | |||
| func TestRelayCapture_RedactsSensitiveRequestHeaders(t *testing.T) { | |||
| clearCaptureCache() | |||
| t.Cleanup(func() { os.RemoveAll(captureTestDir) }) | |||
| SetCaptureEnabled(42, true) | |||
| defer SetCaptureEnabled(42, false) | |||
| requestID := fmt.Sprintf("test-redaction-%d", time.Now().UnixNano()) | |||
| router := setupCaptureTestRouter() | |||
| router.Use(func(c *gin.Context) { | |||
| c.Set("id", 42) | |||
| c.Set(common.RequestIdKey, requestID) | |||
| c.Next() | |||
| }) | |||
| router.Use(RelayCaptureMiddleware()) | |||
| router.POST("/v1/chat/completions", func(c *gin.Context) { | |||
| c.JSON(200, gin.H{"result": "ok"}) | |||
| }) | |||
| req, _ := http.NewRequest("POST", "/v1/chat/completions", bytes.NewReader([]byte(`{"model":"gpt-4"}`))) | |||
| req.Header.Set("Authorization", "Bearer sk-test") | |||
| req.Header.Set("X-Api-Key", "x-secret") | |||
| req.Header.Set("Cookie", "session=secret") | |||
| rec := httptest.NewRecorder() | |||
| router.ServeHTTP(rec, req) | |||
| logPath := filepath.Join(dateDir, files[len(files)-1].Name()) | |||
| content, err := os.ReadFile(logPath) | |||
| assert.Equal(t, 200, rec.Code) | |||
| today := time.Now().Format("2006-01-02") | |||
| path := filepath.Join(captureTestDir, today, requestID+".json") | |||
| content, err := os.ReadFile(path) | |||
| require.NoError(t, err) | |||
| contentStr := string(content) | |||
| t.Logf("Stream capture file content:\n%s", contentStr) | |||
| assert.Contains(t, contentStr, "=== REQUEST") | |||
| assert.Contains(t, contentStr, "POST /v1/chat/completions") | |||
| assert.Contains(t, contentStr, "user_id: 42") | |||
| assert.Contains(t, contentStr, "=== RESPONSE") | |||
| assert.Contains(t, contentStr, "Hello ") | |||
| assert.Contains(t, contentStr, "world!") | |||
| assert.Contains(t, contentStr, "[DONE]") | |||
| assert.Contains(t, contentStr, "=== END") | |||
| assert.Contains(t, contentStr, "duration_ms=") | |||
| assert.Contains(t, contentStr, "[REDACTED]") | |||
| assert.NotContains(t, contentStr, "Bearer sk-test") | |||
| assert.NotContains(t, contentStr, "x-secret") | |||
| assert.NotContains(t, contentStr, "session=secret") | |||
| } | |||
| // --- 未开启抓包用户测试 --- | |||
| func TestRelayCapture_DisabledUser(t *testing.T) { | |||
| clearCaptureCache() | |||
| router := setupCaptureTestRouter() | |||
| router.Use(func(c *gin.Context) { | |||
| c.Set("id", 99) | |||
| c.Set(common.RequestIdKey, "test-disabled-"+time.Now().Format("150405")) | |||
| c.Set(common.RequestIdKey, fmt.Sprintf("test-disabled-%d", time.Now().UnixNano())) | |||
| c.Next() | |||
| }) | |||
| router.Use(RelayCaptureMiddleware()) | |||
| @@ -185,8 +263,6 @@ func TestRelayCapture_DisabledUser(t *testing.T) { | |||
| } | |||
| } | |||
| // --- 大量数据流测试 --- | |||
| func TestRelayCapture_LargeStream(t *testing.T) { | |||
| clearCaptureCache() | |||
| t.Cleanup(func() { os.RemoveAll(captureTestDir) }) | |||
| @@ -194,10 +270,11 @@ func TestRelayCapture_LargeStream(t *testing.T) { | |||
| SetCaptureEnabled(42, true) | |||
| defer SetCaptureEnabled(42, false) | |||
| requestID := fmt.Sprintf("test-large-%d", time.Now().UnixNano()) | |||
| router := setupCaptureTestRouter() | |||
| router.Use(func(c *gin.Context) { | |||
| c.Set("id", 42) | |||
| c.Set(common.RequestIdKey, "test-large-"+time.Now().Format("150405")) | |||
| c.Set(common.RequestIdKey, requestID) | |||
| c.Next() | |||
| }) | |||
| router.Use(RelayCaptureMiddleware()) | |||
| @@ -207,10 +284,10 @@ func TestRelayCapture_LargeStream(t *testing.T) { | |||
| for i := 0; i < 50; i++ { | |||
| chunk := fmt.Sprintf("data: {\"id\":\"%d\",\"choices\":[{\"delta\":{\"content\":\"chunk%d \"}}]}\n\n", i, i) | |||
| c.Writer.Write([]byte(chunk)) | |||
| _, _ = c.Writer.Write([]byte(chunk)) | |||
| flusher.Flush() | |||
| } | |||
| c.Writer.Write([]byte("data: [DONE]\n\n")) | |||
| _, _ = c.Writer.Write([]byte("data: [DONE]\n\n")) | |||
| flusher.Flush() | |||
| }) | |||
| @@ -222,24 +299,13 @@ func TestRelayCapture_LargeStream(t *testing.T) { | |||
| assert.Equal(t, 200, rec.Code) | |||
| today := time.Now().Format("2006-01-02") | |||
| dateDir := filepath.Join(captureTestDir, today) | |||
| files, err := os.ReadDir(dateDir) | |||
| require.NoError(t, err) | |||
| require.True(t, len(files) > 0) | |||
| logPath := filepath.Join(dateDir, files[len(files)-1].Name()) | |||
| content, err := os.ReadFile(logPath) | |||
| require.NoError(t, err) | |||
| contentStr := string(content) | |||
| assert.Contains(t, contentStr, "chunk0") | |||
| assert.Contains(t, contentStr, "chunk49") | |||
| assert.Contains(t, contentStr, "[DONE]") | |||
| record := readCaptureRecord(t, requestID) | |||
| assert.Contains(t, record.Response.Body, "chunk0") | |||
| assert.Contains(t, record.Response.Body, "chunk49") | |||
| assert.Contains(t, record.Response.Body, "[DONE]") | |||
| assert.Len(t, record.Response.BodyChunks, 51) | |||
| } | |||
| // --- GET 请求测试 --- | |||
| func TestRelayCapture_GetRequest(t *testing.T) { | |||
| clearCaptureCache() | |||
| t.Cleanup(func() { os.RemoveAll(captureTestDir) }) | |||
| @@ -247,10 +313,11 @@ func TestRelayCapture_GetRequest(t *testing.T) { | |||
| SetCaptureEnabled(42, true) | |||
| defer SetCaptureEnabled(42, false) | |||
| requestID := fmt.Sprintf("test-get-%d", time.Now().UnixNano()) | |||
| router := setupCaptureTestRouter() | |||
| router.Use(func(c *gin.Context) { | |||
| c.Set("id", 42) | |||
| c.Set(common.RequestIdKey, "test-get-"+time.Now().Format("150405")) | |||
| c.Set(common.RequestIdKey, requestID) | |||
| c.Next() | |||
| }) | |||
| router.Use(RelayCaptureMiddleware()) | |||
| @@ -268,32 +335,20 @@ func TestRelayCapture_GetRequest(t *testing.T) { | |||
| assert.Equal(t, 200, rec.Code) | |||
| today := time.Now().Format("2006-01-02") | |||
| dateDir := filepath.Join(captureTestDir, today) | |||
| files, err := os.ReadDir(dateDir) | |||
| require.NoError(t, err) | |||
| require.True(t, len(files) > 0) | |||
| logPath := filepath.Join(dateDir, files[len(files)-1].Name()) | |||
| content, err := os.ReadFile(logPath) | |||
| require.NoError(t, err) | |||
| contentStr := string(content) | |||
| assert.Contains(t, contentStr, "=== REQUEST") | |||
| assert.Contains(t, contentStr, "GET /v1/models") | |||
| assert.Contains(t, contentStr, "=== RESPONSE") | |||
| assert.Contains(t, contentStr, "gpt-4") | |||
| assert.Contains(t, contentStr, "=== END") | |||
| record := readCaptureRecord(t, requestID) | |||
| assert.Equal(t, "GET", record.Request.Method) | |||
| assert.Equal(t, "/v1/models", record.Request.Path) | |||
| assert.Equal(t, "", record.Request.Body) | |||
| assert.Equal(t, 200, record.Response.StatusCode) | |||
| assert.Contains(t, record.Response.Body, "gpt-4") | |||
| } | |||
| // --- 零 ID 用户测试 --- | |||
| func TestRelayCapture_ZeroUserID(t *testing.T) { | |||
| clearCaptureCache() | |||
| router := setupCaptureTestRouter() | |||
| router.Use(func(c *gin.Context) { | |||
| c.Set(common.RequestIdKey, "test-zero-"+time.Now().Format("150405")) | |||
| c.Set(common.RequestIdKey, fmt.Sprintf("test-zero-%d", time.Now().UnixNano())) | |||
| c.Next() | |||
| }) | |||
| router.Use(RelayCaptureMiddleware()) | |||
| @@ -308,7 +363,42 @@ func TestRelayCapture_ZeroUserID(t *testing.T) { | |||
| assert.Equal(t, 200, rec.Code) | |||
| } | |||
| // --- 验证响应不被篡改 --- | |||
| func TestRelayCapture_EmptyRequestIDUsesFallbackFilename(t *testing.T) { | |||
| clearCaptureCache() | |||
| t.Cleanup(func() { os.RemoveAll(captureTestDir) }) | |||
| SetCaptureEnabled(42, true) | |||
| defer SetCaptureEnabled(42, false) | |||
| router := setupCaptureTestRouter() | |||
| router.Use(func(c *gin.Context) { | |||
| c.Set("id", 42) | |||
| c.Next() | |||
| }) | |||
| router.Use(RelayCaptureMiddleware()) | |||
| router.GET("/v1/models", func(c *gin.Context) { | |||
| c.JSON(200, gin.H{"data": []gin.H{}}) | |||
| }) | |||
| req, _ := http.NewRequest("GET", "/v1/models", nil) | |||
| rec := httptest.NewRecorder() | |||
| router.ServeHTTP(rec, req) | |||
| assert.Equal(t, 200, rec.Code) | |||
| today := time.Now().Format("2006-01-02") | |||
| dateDir := filepath.Join(captureTestDir, today) | |||
| files, err := os.ReadDir(dateDir) | |||
| require.NoError(t, err) | |||
| require.Len(t, files, 1) | |||
| assert.True(t, strings.HasPrefix(files[0].Name(), "capture-")) | |||
| assert.True(t, strings.HasSuffix(files[0].Name(), ".json")) | |||
| record := readLatestCaptureRecord(t) | |||
| assert.True(t, strings.HasPrefix(record.RequestID, "capture-")) | |||
| assert.Equal(t, "GET", record.Request.Method) | |||
| assert.Equal(t, "/v1/models", record.Request.Path) | |||
| } | |||
| func TestRelayCapture_ResponseUnchanged(t *testing.T) { | |||
| clearCaptureCache() | |||
| @@ -320,7 +410,7 @@ func TestRelayCapture_ResponseUnchanged(t *testing.T) { | |||
| router := setupCaptureTestRouter() | |||
| router.Use(func(c *gin.Context) { | |||
| c.Set("id", 42) | |||
| c.Set(common.RequestIdKey, "test-unchanged-"+time.Now().Format("150405")) | |||
| c.Set(common.RequestIdKey, fmt.Sprintf("test-unchanged-%d", time.Now().UnixNano())) | |||
| c.Next() | |||
| }) | |||
| router.Use(RelayCaptureMiddleware()) | |||
| @@ -342,8 +432,6 @@ func TestRelayCapture_ResponseUnchanged(t *testing.T) { | |||
| assert.Contains(t, rec.Body.String(), "chatcmpl-123") | |||
| } | |||
| // --- 并发安全测试 --- | |||
| func TestRelayCapture_ConcurrentRequests(t *testing.T) { | |||
| clearCaptureCache() | |||
| t.Cleanup(func() { os.RemoveAll(captureTestDir) }) | |||
| @@ -352,7 +440,7 @@ func TestRelayCapture_ConcurrentRequests(t *testing.T) { | |||
| defer SetCaptureEnabled(42, false) | |||
| router := setupCaptureTestRouter() | |||
| var reqCounter int64 = 0 | |||
| var reqCounter int64 | |||
| router.Use(func(c *gin.Context) { | |||
| c.Set("id", 42) | |||
| n := atomic.AddInt64(&reqCounter, 1) | |||
| @@ -389,14 +477,17 @@ func TestRelayCapture_ConcurrentRequests(t *testing.T) { | |||
| dateDir := filepath.Join(captureTestDir, today) | |||
| files, err := os.ReadDir(dateDir) | |||
| require.NoError(t, err) | |||
| assert.Equal(t, 10, len(files), "should have 10 capture log files") | |||
| assert.Equal(t, 10, len(files), "should have 10 capture json files") | |||
| for _, f := range files { | |||
| require.True(t, strings.HasSuffix(f.Name(), ".json")) | |||
| content, err := os.ReadFile(filepath.Join(dateDir, f.Name())) | |||
| require.NoError(t, err) | |||
| contentStr := string(content) | |||
| assert.Contains(t, contentStr, "=== REQUEST") | |||
| assert.Contains(t, contentStr, "=== RESPONSE") | |||
| assert.Contains(t, contentStr, "=== END") | |||
| var record relayCaptureTestRecord | |||
| require.NoError(t, common.Unmarshal(content, &record)) | |||
| assert.Equal(t, "POST", record.Request.Method) | |||
| assert.Equal(t, "/v1/chat/completions", record.Request.Path) | |||
| assert.Equal(t, 200, record.Response.StatusCode) | |||
| assert.Contains(t, record.Response.Body, "ok") | |||
| } | |||
| } | |||
| @@ -1,13 +1,13 @@ | |||
| package middleware | |||
| import ( | |||
| "fmt" | |||
| "io" | |||
| "bytes" | |||
| "net/http" | |||
| "net/http/httptest" | |||
| "os" | |||
| "path/filepath" | |||
| "strings" | |||
| "testing" | |||
| "time" | |||
| "github.com/gin-gonic/gin" | |||
| "github.com/stretchr/testify/assert" | |||
| @@ -67,17 +67,15 @@ func TestGetCaptureEnabledUsersCount(t *testing.T) { | |||
| SetCaptureEnabled(2, false) | |||
| } | |||
| // --- captureResponseWriter 测试 --- | |||
| func TestCaptureResponseWriter_Write(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| rec := httptest.NewRecorder() | |||
| c, _ := gin.CreateTestContext(rec) | |||
| ch := make(chan []byte, 16) | |||
| collector := newResponseCaptureCollector() | |||
| cw := &captureResponseWriter{ | |||
| ResponseWriter: c.Writer, | |||
| ch: ch, | |||
| collector: collector, | |||
| } | |||
| n, err := cw.Write([]byte("hello")) | |||
| @@ -85,23 +83,19 @@ func TestCaptureResponseWriter_Write(t *testing.T) { | |||
| assert.Equal(t, 5, n) | |||
| assert.Equal(t, "hello", rec.Body.String()) | |||
| select { | |||
| case data := <-ch: | |||
| assert.Equal(t, []byte("hello"), data) | |||
| default: | |||
| t.Fatal("expected data in channel") | |||
| } | |||
| assert.Equal(t, "hello", collector.Body()) | |||
| assert.Equal(t, []string{"hello"}, collector.Chunks()) | |||
| } | |||
| func TestCaptureResponseWriter_WriteString(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| rec := httptest.NewRecorder() | |||
| c, _ := gin.CreateTestContext(rec) | |||
| ch := make(chan []byte, 16) | |||
| collector := newResponseCaptureCollector() | |||
| cw := &captureResponseWriter{ | |||
| ResponseWriter: c.Writer, | |||
| ch: ch, | |||
| collector: collector, | |||
| } | |||
| n, err := cw.WriteString("test-string") | |||
| @@ -109,86 +103,135 @@ func TestCaptureResponseWriter_WriteString(t *testing.T) { | |||
| assert.Equal(t, len("test-string"), n) | |||
| assert.Equal(t, "test-string", rec.Body.String()) | |||
| select { | |||
| case data := <-ch: | |||
| assert.Equal(t, []byte("test-string"), data) | |||
| default: | |||
| t.Fatal("expected data in channel") | |||
| } | |||
| assert.Equal(t, "test-string", collector.Body()) | |||
| assert.Equal(t, []string{"test-string"}, collector.Chunks()) | |||
| } | |||
| // --- startCaptureWriter 测试 --- | |||
| func TestStartCaptureWriter_WritesAllChunks(t *testing.T) { | |||
| f, err := os.CreateTemp("", "test-capture-*.log") | |||
| require.NoError(t, err) | |||
| filePath := f.Name() | |||
| defer os.Remove(filePath) | |||
| type partialResponseWriter struct { | |||
| gin.ResponseWriter | |||
| } | |||
| start := time.Now() | |||
| chIn, doneCh := startCaptureWriter(f, "test-req", start) | |||
| func (w partialResponseWriter) Write(b []byte) (int, error) { | |||
| _, err := w.ResponseWriter.Write(b[:2]) | |||
| return 2, err | |||
| } | |||
| // 发送 130 个 chunk(大于 channel buffer 128) | |||
| for i := 0; i < 130; i++ { | |||
| chIn <- []byte(fmt.Sprintf("data: %d\n", i)) | |||
| } | |||
| close(chIn) | |||
| func TestCaptureResponseWriter_RecordsOnlyWrittenBytes(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| rec := httptest.NewRecorder() | |||
| c, _ := gin.CreateTestContext(rec) | |||
| collector := newResponseCaptureCollector() | |||
| select { | |||
| case <-doneCh: | |||
| case <-time.After(5 * time.Second): | |||
| t.Fatal("timeout waiting for doneCh") | |||
| cw := &captureResponseWriter{ | |||
| ResponseWriter: partialResponseWriter{ResponseWriter: c.Writer}, | |||
| collector: collector, | |||
| } | |||
| content, err := os.ReadFile(filePath) | |||
| require.NoError(t, err) | |||
| contentStr := string(content) | |||
| assert.Contains(t, contentStr, "data: 0") | |||
| assert.Contains(t, contentStr, "data: 19") | |||
| assert.Contains(t, contentStr, "=== END") | |||
| assert.Contains(t, contentStr, "duration_ms=") | |||
| n, err := cw.Write([]byte("hello")) | |||
| assert.NoError(t, err) | |||
| assert.Equal(t, 2, n) | |||
| assert.Equal(t, "he", rec.Body.String()) | |||
| assert.Equal(t, "he", collector.Body()) | |||
| assert.Equal(t, []string{"he"}, collector.Chunks()) | |||
| } | |||
| // --- 辅助函数测试 --- | |||
| func TestCreateCaptureFile(t *testing.T) { | |||
| f, err := createCaptureFile("test-req-123") | |||
| require.NoError(t, err) | |||
| defer os.RemoveAll(filepath.Dir(f.Name())) | |||
| defer f.Close() | |||
| assert.Contains(t, f.Name(), "relay-capture") | |||
| assert.Contains(t, f.Name(), "test-req-123.log") | |||
| assert.Contains(t, f.Name(), "test-req-123.json") | |||
| assert.NotContains(t, f.Name(), ".log") | |||
| dir := filepath.Dir(f.Name()) | |||
| info, err := os.Stat(dir) | |||
| require.NoError(t, err) | |||
| assert.True(t, info.IsDir()) | |||
| } | |||
| func TestNormalizeCaptureRequestID_UsesExistingID(t *testing.T) { | |||
| requestID := normalizeCaptureRequestID("req-123") | |||
| assert.Equal(t, "req-123", requestID) | |||
| } | |||
| f.Close() | |||
| func TestNormalizeCaptureRequestID_FallbackForEmptyID(t *testing.T) { | |||
| requestID := normalizeCaptureRequestID("") | |||
| assert.True(t, strings.HasPrefix(requestID, "capture-")) | |||
| assert.False(t, strings.HasSuffix(requestID, ".json")) | |||
| } | |||
| func TestWriteRequestBlock(t *testing.T) { | |||
| func TestCopyHeadersForCapture_RedactsSensitiveRequestHeaders(t *testing.T) { | |||
| headers := http.Header{ | |||
| "Authorization": []string{"Bearer sk-test"}, | |||
| "authorization": []string{"Bearer lower"}, | |||
| "X-Api-Key": []string{"x-key"}, | |||
| "Api-Key": []string{"api-key"}, | |||
| "Cookie": []string{"session=secret"}, | |||
| "Proxy-Authorization": []string{"Basic secret"}, | |||
| "Content-Type": []string{"application/json"}, | |||
| "X-Trace-ID": []string{"trace-1"}, | |||
| } | |||
| copied := copyHeadersForCapture(headers, true) | |||
| assert.Equal(t, []string{"[REDACTED]"}, copied["Authorization"]) | |||
| assert.Equal(t, []string{"[REDACTED]"}, copied["authorization"]) | |||
| assert.Equal(t, []string{"[REDACTED]"}, copied["X-Api-Key"]) | |||
| assert.Equal(t, []string{"[REDACTED]"}, copied["Api-Key"]) | |||
| assert.Equal(t, []string{"[REDACTED]"}, copied["Cookie"]) | |||
| assert.Equal(t, []string{"[REDACTED]"}, copied["Proxy-Authorization"]) | |||
| assert.Equal(t, []string{"application/json"}, copied["Content-Type"]) | |||
| assert.Equal(t, []string{"trace-1"}, copied["X-Trace-ID"]) | |||
| assert.Equal(t, []string{"Bearer sk-test"}, headers["Authorization"]) | |||
| } | |||
| func TestCopyHeadersForCapture_ResponseHeadersNotRedacted(t *testing.T) { | |||
| headers := http.Header{} | |||
| headers.Add("Authorization", "response-token") | |||
| headers.Add("Content-Type", "application/json") | |||
| copied := copyHeadersForCapture(headers, false) | |||
| assert.Equal(t, []string{"response-token"}, copied["Authorization"]) | |||
| assert.Equal(t, []string{"application/json"}, copied["Content-Type"]) | |||
| } | |||
| func TestBuildCaptureRequestRecord_PostBody(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| w := httptest.NewRecorder() | |||
| c, _ := gin.CreateTestContext(w) | |||
| c.Set("id", int64(42)) | |||
| c.Request = httptest.NewRequest("POST", "/v1/chat/completions", nil) | |||
| body := `{"model":"gpt-4"}` | |||
| c.Request = httptest.NewRequest("POST", "/v1/chat/completions?stream=true", bytes.NewReader([]byte(body))) | |||
| c.Request.Header.Set("Authorization", "Bearer sk-test") | |||
| c.Request.Header.Set("Content-Type", "application/json") | |||
| f, err := os.CreateTemp("", "test-req-block-*.log") | |||
| require.NoError(t, err) | |||
| filePath := f.Name() | |||
| defer os.Remove(filePath) | |||
| var captureErrors []string | |||
| record := buildCaptureRequestRecord(c, &captureErrors) | |||
| writeRequestBlock(c, f, 42) | |||
| assert.Equal(t, "POST", record.Method) | |||
| assert.Equal(t, "/v1/chat/completions", record.Path) | |||
| assert.Equal(t, "stream=true", record.Query) | |||
| assert.Equal(t, body, record.Body) | |||
| assert.Equal(t, []string{"[REDACTED]"}, record.Headers["Authorization"]) | |||
| assert.Equal(t, []string{"application/json"}, record.Headers["Content-Type"]) | |||
| assert.Empty(t, captureErrors) | |||
| } | |||
| f.Seek(0, io.SeekStart) | |||
| content, err := io.ReadAll(f) | |||
| require.NoError(t, err) | |||
| contentStr := string(content) | |||
| assert.Contains(t, contentStr, "=== REQUEST") | |||
| assert.Contains(t, contentStr, "POST /v1/chat/completions") | |||
| assert.Contains(t, contentStr, "user_id: 42") | |||
| func TestBuildCaptureRequestRecord_GetNoBody(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| w := httptest.NewRecorder() | |||
| c, _ := gin.CreateTestContext(w) | |||
| c.Request = httptest.NewRequest("GET", "/v1/models", nil) | |||
| var captureErrors []string | |||
| record := buildCaptureRequestRecord(c, &captureErrors) | |||
| assert.Equal(t, "GET", record.Method) | |||
| assert.Equal(t, "/v1/models", record.Path) | |||
| assert.Equal(t, "", record.Query) | |||
| assert.Equal(t, "", record.Body) | |||
| assert.Empty(t, captureErrors) | |||
| } | |||
| @@ -28,6 +28,8 @@ type AbilityWithChannel struct { | |||
| ChannelType int `json:"channel_type"` | |||
| } | |||
| type ChannelFilter func(*Channel) (bool, error) | |||
| func GetAllEnableAbilityWithChannels() ([]AbilityWithChannel, error) { | |||
| var abilities []AbilityWithChannel | |||
| err := DB.Table("abilities"). | |||
| @@ -66,7 +68,6 @@ func GetAbilitiesByChannelId(channelId int) ([]*Ability, error) { | |||
| return abilities, err | |||
| } | |||
| func getPriority(group string, model string, retry int) (int, error) { | |||
| var priorities []int | |||
| @@ -113,6 +114,10 @@ func getChannelQuery(group string, model string, retry int) (*gorm.DB, error) { | |||
| } | |||
| func GetChannel(group string, model string, retry int) (*Channel, error) { | |||
| return GetChannelWithFilter(group, model, retry, nil) | |||
| } | |||
| func GetChannelWithFilter(group string, model string, retry int, filter ChannelFilter) (*Channel, error) { | |||
| var abilities []Ability | |||
| var err error = nil | |||
| @@ -128,6 +133,40 @@ func GetChannel(group string, model string, retry int) (*Channel, error) { | |||
| if err != nil { | |||
| return nil, err | |||
| } | |||
| if filter != nil && len(abilities) > 0 { | |||
| channelIds := make([]int, 0, len(abilities)) | |||
| seenChannelIds := make(map[int]struct{}, len(abilities)) | |||
| for _, ability := range abilities { | |||
| if _, ok := seenChannelIds[ability.ChannelId]; ok { | |||
| continue | |||
| } | |||
| seenChannelIds[ability.ChannelId] = struct{}{} | |||
| channelIds = append(channelIds, ability.ChannelId) | |||
| } | |||
| var channels []Channel | |||
| if err = DB.Where("id IN ?", channelIds).Find(&channels).Error; err != nil { | |||
| return nil, err | |||
| } | |||
| channelByID := make(map[int]*Channel, len(channels)) | |||
| for i := range channels { | |||
| channelByID[channels[i].Id] = &channels[i] | |||
| } | |||
| filtered := make([]Ability, 0, len(abilities)) | |||
| for _, ability := range abilities { | |||
| candidate, ok := channelByID[ability.ChannelId] | |||
| if !ok { | |||
| return nil, fmt.Errorf("channel #%d does not exist", ability.ChannelId) | |||
| } | |||
| accepted, filterErr := filter(candidate) | |||
| if filterErr != nil { | |||
| return nil, filterErr | |||
| } | |||
| if accepted { | |||
| filtered = append(filtered, ability) | |||
| } | |||
| } | |||
| abilities = filtered | |||
| } | |||
| channel := Channel{} | |||
| if len(abilities) > 0 { | |||
| // Randomly choose one | |||
| @@ -94,9 +94,13 @@ func SyncChannelCache(frequency int) { | |||
| } | |||
| func GetRandomSatisfiedChannel(group string, model string, retry int) (*Channel, error) { | |||
| return GetRandomSatisfiedChannelWithFilter(group, model, retry, nil) | |||
| } | |||
| func GetRandomSatisfiedChannelWithFilter(group string, model string, retry int, filter ChannelFilter) (*Channel, error) { | |||
| // if memory cache is disabled, get channel directly from database | |||
| if !common.MemoryCacheEnabled { | |||
| return GetChannel(group, model, retry) | |||
| return GetChannelWithFilter(group, model, retry, filter) | |||
| } | |||
| channelSyncLock.RLock() | |||
| @@ -117,6 +121,15 @@ func GetRandomSatisfiedChannel(group string, model string, retry int) (*Channel, | |||
| if len(channels) == 1 { | |||
| if channel, ok := channelsIDM[channels[0]]; ok { | |||
| if filter != nil { | |||
| accepted, err := filter(channel) | |||
| if err != nil { | |||
| return nil, err | |||
| } | |||
| if !accepted { | |||
| return nil, nil | |||
| } | |||
| } | |||
| return channel, nil | |||
| } | |||
| return nil, fmt.Errorf("数据库一致性错误,渠道# %d 不存在,请联系管理员修复", channels[0]) | |||
| @@ -147,6 +160,15 @@ func GetRandomSatisfiedChannel(group string, model string, retry int) (*Channel, | |||
| for _, channelId := range channels { | |||
| if channel, ok := channelsIDM[channelId]; ok { | |||
| if channel.GetPriority() == targetPriority { | |||
| if filter != nil { | |||
| accepted, err := filter(channel) | |||
| if err != nil { | |||
| return nil, err | |||
| } | |||
| if !accepted { | |||
| continue | |||
| } | |||
| } | |||
| sumWeight += channel.GetWeight() | |||
| targetChannels = append(targetChannels, channel) | |||
| } | |||
| @@ -156,6 +178,9 @@ func GetRandomSatisfiedChannel(group string, model string, retry int) (*Channel, | |||
| } | |||
| if len(targetChannels) == 0 { | |||
| if filter != nil { | |||
| return nil, nil | |||
| } | |||
| return nil, errors.New(fmt.Sprintf("no channel found, group: %s, model: %s, priority: %d", group, model, targetPriority)) | |||
| } | |||
| @@ -4,6 +4,7 @@ import ( | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/constant" | |||
| "github.com/glebarez/sqlite" | |||
| "github.com/stretchr/testify/assert" | |||
| "github.com/stretchr/testify/require" | |||
| @@ -50,9 +51,9 @@ func setupChannelSelectDB(t *testing.T) *gorm.DB { | |||
| return db | |||
| } | |||
| func int64Ptr(v int64) *int64 { return &v } | |||
| func intPtr(v int) *int { return &v } | |||
| func uintPtr(v uint) *uint { return &v } | |||
| func int64Ptr(v int64) *int64 { return &v } | |||
| func intPtr(v int) *int { return &v } | |||
| func uintPtr(v uint) *uint { return &v } | |||
| func createChannelForTest(t *testing.T, db *gorm.DB, id int, name, group, models string, priority int64) { | |||
| t.Helper() | |||
| @@ -72,6 +73,24 @@ func createChannelForTest(t *testing.T, db *gorm.DB, id int, name, group, models | |||
| require.NoError(t, db.Create(&ch).Error) | |||
| } | |||
| func createChannelWithTypeForTest(t *testing.T, db *gorm.DB, id int, name, group, models string, channelType int, priority int64) { | |||
| t.Helper() | |||
| ch := Channel{ | |||
| Id: id, | |||
| Type: channelType, | |||
| Key: "test-key-" + name, | |||
| Status: common.ChannelStatusEnabled, | |||
| Name: name, | |||
| Group: group, | |||
| Models: models, | |||
| Priority: int64Ptr(priority), | |||
| Weight: uintPtr(10), | |||
| AutoBan: intPtr(1), | |||
| CreatedTime: 1, | |||
| } | |||
| require.NoError(t, db.Create(&ch).Error) | |||
| } | |||
| func createAbilityForTest(t *testing.T, db *gorm.DB, group, modelName string, channelId int, priority int64, enabled bool) { | |||
| t.Helper() | |||
| ability := Ability{ | |||
| @@ -185,3 +204,36 @@ func TestGetChannel_MultipleChannelsWeighted(t *testing.T) { | |||
| assert.True(t, found[1], "channel 1 should be selected at least once") | |||
| assert.True(t, found[2], "channel 2 should be selected at least once") | |||
| } | |||
| func TestGetRandomSatisfiedChannelWithFilter_AllCandidatesRejected(t *testing.T) { | |||
| db := setupChannelSelectDB(t) | |||
| createChannelForTest(t, db, 1, "ch1", "default", "seedance-2", 10) | |||
| createChannelForTest(t, db, 2, "ch2", "default", "seedance-2", 10) | |||
| createAbilityForTest(t, db, "default", "seedance-2", 1, 10, true) | |||
| createAbilityForTest(t, db, "default", "seedance-2", 2, 10, true) | |||
| ch, err := GetRandomSatisfiedChannelWithFilter("default", "seedance-2", 0, func(*Channel) (bool, error) { | |||
| return false, nil | |||
| }) | |||
| require.NoError(t, err) | |||
| assert.Nil(t, ch) | |||
| } | |||
| func TestGetRandomSatisfiedChannelWithFilter_SelectsOnlyAcceptedCandidate(t *testing.T) { | |||
| db := setupChannelSelectDB(t) | |||
| createChannelWithTypeForTest(t, db, 1, "openai", "default", "seedance-2", constant.ChannelTypeOpenAI, 10) | |||
| createChannelWithTypeForTest(t, db, 2, "doubao", "default", "seedance-2", constant.ChannelTypeDoubaoVideo, 10) | |||
| createAbilityForTest(t, db, "default", "seedance-2", 1, 10, true) | |||
| createAbilityForTest(t, db, "default", "seedance-2", 2, 10, true) | |||
| for i := 0; i < 20; i++ { | |||
| ch, err := GetRandomSatisfiedChannelWithFilter("default", "seedance-2", 0, func(channel *Channel) (bool, error) { | |||
| return channel.Type == constant.ChannelTypeDoubaoVideo, nil | |||
| }) | |||
| require.NoError(t, err) | |||
| require.NotNil(t, ch) | |||
| assert.Equal(t, 2, ch.Id) | |||
| } | |||
| } | |||
| @@ -211,15 +211,16 @@ func RecordConsumeLog(c *gin.Context, userId int, params RecordConsumeLogParams) | |||
| } | |||
| type RecordTaskBillingLogParams struct { | |||
| UserId int | |||
| LogType int | |||
| Content string | |||
| ChannelId int | |||
| ModelName string | |||
| Quota int | |||
| TokenId int | |||
| Group string | |||
| Other map[string]interface{} | |||
| UserId int | |||
| LogType int | |||
| Content string | |||
| ChannelId int | |||
| ModelName string | |||
| Quota int | |||
| TokenId int | |||
| Group string | |||
| CompletionTokens int | |||
| Other map[string]interface{} | |||
| } | |||
| func RecordTaskBillingLog(params RecordTaskBillingLogParams) { | |||
| @@ -234,18 +235,19 @@ func RecordTaskBillingLog(params RecordTaskBillingLogParams) { | |||
| } | |||
| } | |||
| log := &Log{ | |||
| UserId: params.UserId, | |||
| Username: username, | |||
| CreatedAt: common.GetTimestamp(), | |||
| Type: params.LogType, | |||
| Content: params.Content, | |||
| TokenName: tokenName, | |||
| ModelName: params.ModelName, | |||
| Quota: params.Quota, | |||
| ChannelId: params.ChannelId, | |||
| TokenId: params.TokenId, | |||
| Group: params.Group, | |||
| Other: common.MapToJsonStr(params.Other), | |||
| UserId: params.UserId, | |||
| Username: username, | |||
| CreatedAt: common.GetTimestamp(), | |||
| Type: params.LogType, | |||
| Content: params.Content, | |||
| TokenName: tokenName, | |||
| ModelName: params.ModelName, | |||
| Quota: params.Quota, | |||
| ChannelId: params.ChannelId, | |||
| TokenId: params.TokenId, | |||
| Group: params.Group, | |||
| CompletionTokens: params.CompletionTokens, | |||
| Other: common.MapToJsonStr(params.Other), | |||
| } | |||
| err := LOG_DB.Create(log).Error | |||
| if err != nil { | |||
| @@ -286,6 +286,7 @@ func migrateDB() error { | |||
| &QuotaSyncLog{}, | |||
| &EmailQuotaRule{}, | |||
| &UserModelRateLimit{}, | |||
| &UserAssetChannel{}, | |||
| &UserMigrationBatch{}, | |||
| &UserMigrationItem{}, | |||
| &MigrationQuotaGrant{}, | |||
| @@ -353,6 +354,7 @@ func migrateDBFast() error { | |||
| {&PendingSyncRecord{}, "PendingSyncRecord"}, | |||
| {&QuotaSyncLog{}, "QuotaSyncLog"}, | |||
| {&EmailQuotaRule{}, "EmailQuotaRule"}, | |||
| {&UserAssetChannel{}, "UserAssetChannel"}, | |||
| {&UserMigrationBatch{}, "UserMigrationBatch"}, | |||
| {&UserMigrationItem{}, "UserMigrationItem"}, | |||
| {&MigrationQuotaGrant{}, "MigrationQuotaGrant"}, | |||
| @@ -133,6 +133,8 @@ func InitOptionMap() { | |||
| common.OptionMap["ModelRequestRateLimitGroup"] = setting.ModelRequestRateLimitGroup2JSONString() | |||
| common.OptionMap["ModelRatio"] = ratio_setting.ModelRatio2JSONString() | |||
| common.OptionMap["ModelPrice"] = ratio_setting.ModelPrice2JSONString() | |||
| common.OptionMap["ModelPricingRules"] = ratio_setting.ModelPricingRules2JSONString() | |||
| common.OptionMap["ModelDisplayPricing"] = ratio_setting.ModelDisplayPricing2JSONString() | |||
| common.OptionMap["CacheRatio"] = ratio_setting.CacheRatio2JSONString() | |||
| common.OptionMap["CreateCacheRatio"] = ratio_setting.CreateCacheRatio2JSONString() | |||
| common.OptionMap["GroupRatio"] = ratio_setting.GroupRatio2JSONString() | |||
| @@ -511,6 +513,10 @@ func updateOptionMap(key string, value string) (err error) { | |||
| err = ratio_setting.UpdateCompletionRatioByJSONString(value) | |||
| case "ModelPrice": | |||
| err = ratio_setting.UpdateModelPriceByJSONString(value) | |||
| case "ModelPricingRules": | |||
| err = ratio_setting.UpdateModelPricingRulesByJSONString(value) | |||
| case "ModelDisplayPricing": | |||
| err = ratio_setting.UpdateModelDisplayPricingByJSONString(value) | |||
| case "CacheRatio": | |||
| err = ratio_setting.UpdateCacheRatioByJSONString(value) | |||
| case "CreateCacheRatio": | |||
| @@ -14,22 +14,24 @@ import ( | |||
| ) | |||
| type Pricing struct { | |||
| ModelName string `json:"model_name"` | |||
| Description string `json:"description,omitempty"` | |||
| Icon string `json:"icon,omitempty"` | |||
| Tags string `json:"tags,omitempty"` | |||
| VendorID int `json:"vendor_id,omitempty"` | |||
| QuotaType int `json:"quota_type"` | |||
| ModelRatio float64 `json:"model_ratio"` | |||
| ModelPrice float64 `json:"model_price"` | |||
| OwnerBy string `json:"owner_by"` | |||
| CompletionRatio float64 `json:"completion_ratio"` | |||
| CacheRatio float64 `json:"cache_ratio"` | |||
| CacheCreationRatio float64 `json:"cache_creation_ratio"` | |||
| EnableGroup []string `json:"enable_groups"` | |||
| SupportedEndpointTypes []constant.EndpointType `json:"supported_endpoint_types"` | |||
| PricingVersion string `json:"pricing_version,omitempty"` | |||
| Type int `json:"type"` | |||
| ModelName string `json:"model_name"` | |||
| Description string `json:"description,omitempty"` | |||
| Icon string `json:"icon,omitempty"` | |||
| Tags string `json:"tags,omitempty"` | |||
| VendorID int `json:"vendor_id,omitempty"` | |||
| QuotaType int `json:"quota_type"` | |||
| ModelRatio float64 `json:"model_ratio"` | |||
| ModelPrice float64 `json:"model_price"` | |||
| OwnerBy string `json:"owner_by"` | |||
| CompletionRatio float64 `json:"completion_ratio"` | |||
| CacheRatio float64 `json:"cache_ratio"` | |||
| CacheCreationRatio float64 `json:"cache_creation_ratio"` | |||
| EnableGroup []string `json:"enable_groups"` | |||
| SupportedEndpointTypes []constant.EndpointType `json:"supported_endpoint_types"` | |||
| PricingVersion string `json:"pricing_version,omitempty"` | |||
| Type int `json:"type"` | |||
| PricingConfig *types.PricingConfig `json:"pricing_config,omitempty"` | |||
| DisplayPricing []types.ModelDisplayPricingItem `json:"display_pricing,omitempty"` | |||
| } | |||
| type PricingVendor struct { | |||
| @@ -339,6 +341,7 @@ func updatePricing() { | |||
| } | |||
| applyGlobalDefault(&pricing, model) | |||
| applyDisplayPricing(&pricing, model) | |||
| pricingMap = append(pricingMap, pricing) | |||
| } | |||
| @@ -384,6 +387,9 @@ func GetSupportedEndpointMap() map[string]common.EndpointInfo { | |||
| } | |||
| func applyGlobalDefault(pricing *Pricing, model string) { | |||
| if applyMatrixPricingConfig(pricing, model) { | |||
| return | |||
| } | |||
| modelPrice, findPrice := ratio_setting.GetModelPrice(model, false) | |||
| if findPrice { | |||
| pricing.ModelPrice = modelPrice | |||
| @@ -400,3 +406,67 @@ func applyGlobalDefault(pricing *Pricing, model string) { | |||
| pricing.CacheRatio, _ = ratio_setting.GetCacheRatio(model) | |||
| pricing.CacheCreationRatio, _ = ratio_setting.GetCreateCacheRatio(model) | |||
| } | |||
| func applyMatrixPricingConfig(pricing *Pricing, model string) bool { | |||
| cfg := ratio_setting.GetPricingConfig(model) | |||
| if cfg == nil { | |||
| return false | |||
| } | |||
| pricing.PricingConfig = cfg | |||
| minPrice := minPricingConfigPrice(cfg) | |||
| if cfg.BillingUnit == types.BillingUnitPer1MTokens { | |||
| pricing.ModelPrice = minPrice | |||
| pricing.ModelRatio = 0 | |||
| pricing.CompletionRatio = 0 | |||
| pricing.QuotaType = QuotaTypeByTokens | |||
| return true | |||
| } | |||
| pricing.ModelPrice = minPrice | |||
| pricing.ModelRatio = 0 | |||
| pricing.CompletionRatio = 0 | |||
| pricing.QuotaType = QuotaTypeByCall | |||
| return true | |||
| } | |||
| func applyDisplayPricing(pricing *Pricing, model string) { | |||
| items := ratio_setting.GetModelDisplayPricing(model) | |||
| if len(items) > 0 { | |||
| pricing.DisplayPricing = items | |||
| } | |||
| } | |||
| func minPricingConfigPrice(cfg *types.PricingConfig) float64 { | |||
| if cfg == nil || len(cfg.Table) == 0 { | |||
| return 0 | |||
| } | |||
| var min float64 | |||
| found := false | |||
| for _, row := range cfg.Table { | |||
| price, ok := pricingRowNumber(row["price"]) | |||
| if !ok { | |||
| continue | |||
| } | |||
| if !found || price < min { | |||
| min = price | |||
| found = true | |||
| } | |||
| } | |||
| return min | |||
| } | |||
| func pricingRowNumber(v any) (float64, bool) { | |||
| switch n := v.(type) { | |||
| case float64: | |||
| return n, true | |||
| case float32: | |||
| return float64(n), true | |||
| case int: | |||
| return float64(n), true | |||
| case int64: | |||
| return float64(n), true | |||
| case int32: | |||
| return float64(n), true | |||
| default: | |||
| return 0, false | |||
| } | |||
| } | |||
| @@ -7,6 +7,7 @@ import ( | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/constant" | |||
| "github.com/QuantumNous/new-api/setting/ratio_setting" | |||
| "github.com/QuantumNous/new-api/types" | |||
| "github.com/glebarez/sqlite" | |||
| "github.com/stretchr/testify/require" | |||
| "gorm.io/gorm" | |||
| @@ -18,6 +19,7 @@ const testPricingGlobalModel = "test-pricing-global-model" | |||
| func resetPricingRatios(t *testing.T) { | |||
| t.Helper() | |||
| require.NoError(t, ratio_setting.UpdateModelPricingRulesByJSONString(`{}`)) | |||
| require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{}`)) | |||
| require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(`{}`)) | |||
| require.NoError(t, ratio_setting.UpdateCacheRatioByJSONString(`{}`)) | |||
| @@ -142,6 +144,86 @@ func TestGetPricingByModel_UsesGlobalDefaultOnly(t *testing.T) { | |||
| require.Equal(t, 0.75, pricing.CacheCreationRatio) | |||
| } | |||
| func TestGetPricingByModel_IncludesMatrixPricingConfig(t *testing.T) { | |||
| setupPricingTest(t) | |||
| const matrixModel = "doubao-seedance-2-0-260128" | |||
| require.NoError(t, ratio_setting.SetPricingConfig(matrixModel, &types.PricingConfig{ | |||
| SchemaVersion: types.PricingSchemaVersion, | |||
| Scope: types.PricingScopeModel, | |||
| BillingUnit: types.BillingUnitPer1MTokens, | |||
| PreconsumeStrategy: types.PreconsumeStrategyMinimum, | |||
| Dimensions: []types.PricingDimension{ | |||
| {Key: "resolution", Source: "request.resolution", Type: "string"}, | |||
| {Key: "video_input", Source: "derived.video_input", Type: "boolean"}, | |||
| }, | |||
| Table: []types.PricingRow{ | |||
| {"resolution": "480p", "video_input": true, "price": 28.0, "source": types.PricingRowSourceManual}, | |||
| {"resolution": "480p", "video_input": false, "price": 46.0, "source": types.PricingRowSourceManual}, | |||
| }, | |||
| Fallback: types.PricingFallback{Strategy: types.PricingFallbackReject}, | |||
| })) | |||
| require.NoError(t, (&Model{ | |||
| ModelName: matrixModel, | |||
| Status: 1, | |||
| }).Insert()) | |||
| require.NoError(t, (&Channel{ | |||
| Id: 1, | |||
| Type: constant.ChannelTypeDoubaoVideoCompatibleAiping, | |||
| Key: "k", | |||
| Name: "channel-a", | |||
| Status: common.ChannelStatusEnabled, | |||
| Models: matrixModel, | |||
| Group: "default", | |||
| }).Insert()) | |||
| RefreshPricing() | |||
| pricing := GetPricingByModel(matrixModel) | |||
| require.NotNil(t, pricing) | |||
| require.NotNil(t, pricing.PricingConfig) | |||
| require.Equal(t, types.BillingUnitPer1MTokens, pricing.PricingConfig.BillingUnit) | |||
| require.Equal(t, QuotaTypeByTokens, pricing.QuotaType) | |||
| require.Equal(t, 28.0, pricing.ModelPrice) | |||
| require.Equal(t, []constant.EndpointType{constant.EndpointTypeDoubaoVideo}, pricing.SupportedEndpointTypes) | |||
| } | |||
| func TestGetPricingByModelIncludesDisplayPricing(t *testing.T) { | |||
| setupPricingTest(t) | |||
| const displayModel = "display-pricing-model" | |||
| require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"`+displayModel+`":10}`)) | |||
| require.NoError(t, ratio_setting.SetModelDisplayPricing(displayModel, []types.ModelDisplayPricingItem{ | |||
| {Specification: "768p-6s", Price: 0.33333, Unit: "second", SortOrder: 1}, | |||
| })) | |||
| require.NoError(t, (&Model{ | |||
| ModelName: displayModel, | |||
| Status: 1, | |||
| }).Insert()) | |||
| require.NoError(t, (&Channel{ | |||
| Id: 1, | |||
| Type: constant.ChannelTypeOpenAI, | |||
| Key: "k", | |||
| Name: "channel-a", | |||
| Status: common.ChannelStatusEnabled, | |||
| Models: displayModel, | |||
| Group: "default", | |||
| }).Insert()) | |||
| RefreshPricing() | |||
| pricing := GetPricingByModel(displayModel) | |||
| require.NotNil(t, pricing) | |||
| require.Len(t, pricing.DisplayPricing, 1) | |||
| require.Equal(t, "768p-6s", pricing.DisplayPricing[0].Specification) | |||
| require.Equal(t, QuotaTypeByTokens, pricing.QuotaType) | |||
| require.Equal(t, 10.0, pricing.ModelRatio) | |||
| require.Equal(t, 0.0, pricing.ModelPrice) | |||
| require.Nil(t, pricing.PricingConfig) | |||
| } | |||
| func TestParseModelEndpointTypes_SupportsObjectAndArray(t *testing.T) { | |||
| objectEndpoints := parseModelEndpointTypes(`{ | |||
| "anthropic": {"path": "/v1/messages", "method": "POST"}, | |||
| @@ -42,27 +42,26 @@ const ( | |||
| ) | |||
| type Task struct { | |||
| ID int64 `json:"id" gorm:"primary_key;AUTO_INCREMENT"` | |||
| CreatedAt int64 `json:"created_at" gorm:"index"` | |||
| UpdatedAt int64 `json:"updated_at"` | |||
| TaskID string `json:"task_id" gorm:"type:varchar(191);index"` // 第三方id,不一定有/ song id\ Task id | |||
| Platform constant.TaskPlatform `json:"platform" gorm:"type:varchar(30);index"` // 平台 | |||
| UserId int `json:"user_id" gorm:"index"` | |||
| Group string `json:"group" gorm:"type:varchar(50)"` // 修正计费用 | |||
| ChannelId int `json:"channel_id" gorm:"index"` | |||
| Quota int `json:"quota"` | |||
| Action string `json:"action" gorm:"type:varchar(40);index"` // 任务类型, song, lyrics, description-mode | |||
| Status TaskStatus `json:"status" gorm:"type:varchar(20);index"` // 任务状态 | |||
| FailReason string `json:"fail_reason"` | |||
| SubmitTime int64 `json:"submit_time" gorm:"index"` | |||
| StartTime int64 `json:"start_time" gorm:"index"` | |||
| FinishTime int64 `json:"finish_time" gorm:"index"` | |||
| Progress string `json:"progress" gorm:"type:varchar(20);index"` | |||
| Properties Properties `json:"properties" gorm:"type:json"` | |||
| Username string `json:"username,omitempty" gorm:"-"` | |||
| // 禁止返回给用户,内部可能包含key等隐私信息 | |||
| PrivateData TaskPrivateData `json:"-" gorm:"column:private_data;type:json"` | |||
| Data json.RawMessage `json:"data" gorm:"type:json"` | |||
| ID int64 `json:"id" gorm:"primary_key;AUTO_INCREMENT"` | |||
| CreatedAt int64 `json:"created_at" gorm:"index"` | |||
| UpdatedAt int64 `json:"updated_at"` | |||
| TaskID string `json:"task_id" gorm:"type:varchar(191);index"` | |||
| Platform constant.TaskPlatform `json:"platform" gorm:"type:varchar(30);index"` | |||
| UserId int `json:"user_id" gorm:"index"` | |||
| Group string `json:"group" gorm:"type:varchar(50)"` | |||
| ChannelId int `json:"channel_id" gorm:"index"` | |||
| Quota int `json:"quota"` | |||
| Action string `json:"action" gorm:"type:varchar(40);index"` | |||
| Status TaskStatus `json:"status" gorm:"type:varchar(20);index"` | |||
| FailReason string `json:"fail_reason"` | |||
| SubmitTime int64 `json:"submit_time" gorm:"index"` | |||
| StartTime int64 `json:"start_time" gorm:"index"` | |||
| FinishTime int64 `json:"finish_time" gorm:"index"` | |||
| Progress string `json:"progress" gorm:"type:varchar(20);index"` | |||
| Properties Properties `json:"properties" gorm:"type:json"` | |||
| Username string `json:"username,omitempty" gorm:"-"` | |||
| PrivateData TaskPrivateData `json:"-" gorm:"column:private_data;type:json"` | |||
| Data json.RawMessage `json:"data" gorm:"type:json"` | |||
| } | |||
| func (t *Task) SetData(data any) { | |||
| @@ -97,28 +96,38 @@ func (m Properties) Value() (driver.Value, error) { | |||
| } | |||
| type TaskPrivateData struct { | |||
| Key string `json:"key,omitempty"` | |||
| UpstreamTaskID string `json:"upstream_task_id,omitempty"` // 上游真实 task ID | |||
| ResultURL string `json:"result_url,omitempty"` // 任务成功后的结果 URL(视频地址等) | |||
| // 计费上下文:用于异步退款/差额结算(轮询阶段读取) | |||
| BillingSource string `json:"billing_source,omitempty"` // "wallet" 或 "subscription" | |||
| SubscriptionId int `json:"subscription_id,omitempty"` // 订阅 ID,用于订阅退款 | |||
| TokenId int `json:"token_id,omitempty"` // 令牌 ID,用于令牌额度退款 | |||
| BillingContext *TaskBillingContext `json:"billing_context,omitempty"` // 计费参数快照(用于轮询阶段重新计算) | |||
| Key string `json:"key,omitempty"` | |||
| UpstreamTaskID string `json:"upstream_task_id,omitempty"` | |||
| ResultURL string `json:"result_url,omitempty"` | |||
| BillingSource string `json:"billing_source,omitempty"` | |||
| SubscriptionId int `json:"subscription_id,omitempty"` | |||
| TokenId int `json:"token_id,omitempty"` | |||
| BillingContext *TaskBillingContext `json:"billing_context,omitempty"` | |||
| UpstreamRequest *TaskUpstreamRequestSnapshot `json:"upstream_request,omitempty"` | |||
| } | |||
| type TaskUpstreamRequestSnapshot struct { | |||
| Body any `json:"body,omitempty"` | |||
| RawBytes int `json:"raw_bytes,omitempty"` | |||
| StoredBytes int `json:"stored_bytes,omitempty"` | |||
| Truncated bool `json:"truncated,omitempty"` | |||
| Redacted bool `json:"redacted,omitempty"` | |||
| SHA256 string `json:"sha256,omitempty"` | |||
| } | |||
| // TaskBillingContext 记录任务提交时的计费参数,以便轮询阶段可以重新计算额度。 | |||
| type TaskBillingContext struct { | |||
| ModelPrice float64 `json:"model_price,omitempty"` // 模型单价 | |||
| GroupRatio float64 `json:"group_ratio,omitempty"` // 分组倍率 | |||
| ModelRatio float64 `json:"model_ratio,omitempty"` // 模型倍率 | |||
| OtherRatios map[string]float64 `json:"other_ratios,omitempty"` // 附加倍率(时长、分辨率等) | |||
| OriginModelName string `json:"origin_model_name,omitempty"` // 模型名称,必须为OriginModelName | |||
| PerCallBilling bool `json:"per_call_billing,omitempty"` // 按次计费:跳过轮询阶段的差额结算 | |||
| ModelPrice float64 `json:"model_price,omitempty"` | |||
| GroupRatio float64 `json:"group_ratio,omitempty"` | |||
| ModelRatio float64 `json:"model_ratio,omitempty"` | |||
| OtherRatios map[string]float64 `json:"other_ratios,omitempty"` | |||
| OriginModelName string `json:"origin_model_name,omitempty"` | |||
| PerCallBilling bool `json:"per_call_billing,omitempty"` | |||
| BillingMode string `json:"billing_mode,omitempty"` | |||
| PricingSnapshot map[string]any `json:"pricing_snapshot,omitempty"` | |||
| BillingUnit string `json:"billing_unit,omitempty"` | |||
| TokenUnitPriceUSD float64 `json:"token_unit_price_usd,omitempty"` | |||
| } | |||
| // GetUpstreamTaskID 获取上游真实 task ID(用于与 provider 通信) | |||
| // 旧数据没有 UpstreamTaskID 时,TaskID 本身就是上游 ID | |||
| func (t *Task) GetUpstreamTaskID() string { | |||
| if t.PrivateData.UpstreamTaskID != "" { | |||
| return t.PrivateData.UpstreamTaskID | |||
| @@ -126,8 +135,6 @@ func (t *Task) GetUpstreamTaskID() string { | |||
| return t.TaskID | |||
| } | |||
| // GetResultURL 获取任务结果 URL(视频地址等) | |||
| // 新数据存在 PrivateData.ResultURL 中;旧数据回退到 FailReason(历史兼容) | |||
| func (t *Task) GetResultURL() string { | |||
| if t.PrivateData.ResultURL != "" { | |||
| return t.PrivateData.ResultURL | |||
| @@ -135,7 +142,6 @@ func (t *Task) GetResultURL() string { | |||
| return t.FailReason | |||
| } | |||
| // GenerateTaskID 生成对外暴露的 task_xxxx 格式 ID | |||
| func GenerateTaskID() string { | |||
| key, _ := common.GenerateRandomCharsKey(32) | |||
| return "task_" + key | |||
| @@ -156,9 +162,9 @@ func (p TaskPrivateData) Value() (driver.Value, error) { | |||
| return common.Marshal(p) | |||
| } | |||
| // SyncTaskQueryParams 用于包含所有搜索条件的结构体,可以根据需求添加更多字段 | |||
| type SyncTaskQueryParams struct { | |||
| Platform constant.TaskPlatform | |||
| UserId int `json:"user_id" gorm:"index"` | |||
| ChannelID string | |||
| TaskID string | |||
| UserID string | |||
| @@ -184,7 +190,6 @@ func InitTask(platform constant.TaskPlatform, relayInfo *commonRelay.RelayInfo) | |||
| } | |||
| } | |||
| // 使用预生成的公开 ID(如果有),否则新生成 | |||
| taskID := "" | |||
| if relayInfo.TaskRelayInfo != nil && relayInfo.TaskRelayInfo.PublicTaskID != "" { | |||
| taskID = relayInfo.TaskRelayInfo.PublicTaskID | |||
| @@ -211,7 +216,6 @@ func TaskGetAllUserTask(userId int, startIdx int, num int, queryParams SyncTaskQ | |||
| var tasks []*Task | |||
| var err error | |||
| // 初始化查询构建器 | |||
| query := DB.Where("user_id = ?", userId) | |||
| if queryParams.TaskID != "" { | |||
| @@ -227,14 +231,12 @@ func TaskGetAllUserTask(userId int, startIdx int, num int, queryParams SyncTaskQ | |||
| query = query.Where("platform = ?", queryParams.Platform) | |||
| } | |||
| if queryParams.StartTimestamp != 0 { | |||
| // 假设您已将前端传来的时间戳转换为数据库所需的时间格式,并处理了时间戳的验证和解析 | |||
| query = query.Where("submit_time >= ?", queryParams.StartTimestamp) | |||
| } | |||
| if queryParams.EndTimestamp != 0 { | |||
| query = query.Where("submit_time <= ?", queryParams.EndTimestamp) | |||
| } | |||
| // 获取数据 | |||
| err = query.Omit("channel_id").Order("id desc").Limit(num).Offset(startIdx).Find(&tasks).Error | |||
| if err != nil { | |||
| return nil | |||
| @@ -247,10 +249,8 @@ func TaskGetAllTasks(startIdx int, num int, queryParams SyncTaskQueryParams) []* | |||
| var tasks []*Task | |||
| var err error | |||
| // 初始化查询构建器 | |||
| query := DB | |||
| // 添加过滤条件 | |||
| if queryParams.ChannelID != "" { | |||
| query = query.Where("channel_id = ?", queryParams.ChannelID) | |||
| } | |||
| @@ -279,7 +279,6 @@ func TaskGetAllTasks(startIdx int, num int, queryParams SyncTaskQueryParams) []* | |||
| query = query.Where("submit_time <= ?", queryParams.EndTimestamp) | |||
| } | |||
| // 获取数据 | |||
| err = query.Order("id desc").Limit(num).Offset(startIdx).Find(&tasks).Error | |||
| if err != nil { | |||
| return nil | |||
| @@ -416,7 +415,6 @@ func (t *Task) UpdateWithStatus(fromStatus TaskStatus) (bool, error) { | |||
| } | |||
| // TaskBulkUpdateByID performs an unconditional bulk UPDATE by primary key IDs. | |||
| // WARNING: This function has NO CAS (Compare-And-Swap) guard — it will overwrite | |||
| // any concurrent status changes. DO NOT use in billing/quota lifecycle flows | |||
| // (e.g., timeout, success, failure transitions that trigger refunds or settlements). | |||
| // For status transitions that involve billing, use Task.UpdateWithStatus() instead. | |||
| @@ -122,6 +122,29 @@ func TestSnapshot_Roundtrip(t *testing.T) { | |||
| assert.JSONEq(t, string(task.Data), string(snap.Data)) | |||
| } | |||
| func TestTaskPrivateDataStoresUpstreamRequestSnapshot(t *testing.T) { | |||
| privateData := TaskPrivateData{ | |||
| UpstreamRequest: &TaskUpstreamRequestSnapshot{ | |||
| Body: map[string]any{"model": "doubao-seedance-2-0-260128"}, | |||
| RawBytes: 42, | |||
| StoredBytes: 128, | |||
| Redacted: true, | |||
| SHA256: "abc123", | |||
| }, | |||
| } | |||
| value, err := privateData.Value() | |||
| require.NoError(t, err) | |||
| var restored TaskPrivateData | |||
| require.NoError(t, restored.Scan(value)) | |||
| require.NotNil(t, restored.UpstreamRequest) | |||
| assert.Equal(t, 42, restored.UpstreamRequest.RawBytes) | |||
| assert.Equal(t, 128, restored.UpstreamRequest.StoredBytes) | |||
| assert.True(t, restored.UpstreamRequest.Redacted) | |||
| assert.Equal(t, "abc123", restored.UpstreamRequest.SHA256) | |||
| } | |||
| // --------------------------------------------------------------------------- | |||
| // UpdateWithStatus CAS — DB integration tests | |||
| // --------------------------------------------------------------------------- | |||
| @@ -0,0 +1,87 @@ | |||
| package model | |||
| import ( | |||
| "errors" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "gorm.io/gorm" | |||
| "gorm.io/gorm/clause" | |||
| ) | |||
| type UserAssetChannel struct { | |||
| Id int `json:"id" gorm:"primaryKey"` | |||
| UserId int `json:"user_id" gorm:"not null;uniqueIndex:idx_user_asset_channel,priority:1"` | |||
| ChannelType int `json:"channel_type" gorm:"not null;uniqueIndex:idx_user_asset_channel,priority:2"` | |||
| Group string `json:"group" gorm:"column:group;type:varchar(64);not null;uniqueIndex:idx_user_asset_channel,priority:3"` | |||
| ChannelId int `json:"channel_id" gorm:"not null"` | |||
| CreatedAt int64 `json:"created_at"` | |||
| UpdatedAt int64 `json:"updated_at"` | |||
| } | |||
| func (UserAssetChannel) TableName() string { | |||
| return "user_asset_channels" | |||
| } | |||
| func GetUserAssetChannel(userId int, channelType int, group string) (*UserAssetChannel, error) { | |||
| var binding UserAssetChannel | |||
| err := DB.Where("user_id = ? AND channel_type = ? AND "+commonGroupCol+" = ?", userId, channelType, group).First(&binding).Error | |||
| if errors.Is(err, gorm.ErrRecordNotFound) { | |||
| return nil, nil | |||
| } | |||
| if err != nil { | |||
| return nil, err | |||
| } | |||
| return &binding, nil | |||
| } | |||
| func BindUserAssetChannel(userId int, channelType int, group string, channelId int) error { | |||
| return BindUserAssetChannelWithTx(DB, userId, channelType, group, channelId) | |||
| } | |||
| func BindUserAssetChannelWithTx(tx *gorm.DB, userId int, channelType int, group string, channelId int) error { | |||
| now := common.GetTimestamp() | |||
| binding := UserAssetChannel{ | |||
| UserId: userId, | |||
| ChannelType: channelType, | |||
| Group: group, | |||
| ChannelId: channelId, | |||
| CreatedAt: now, | |||
| UpdatedAt: now, | |||
| } | |||
| return tx.Session(&gorm.Session{SkipDefaultTransaction: true}).Clauses(clause.OnConflict{ | |||
| Columns: []clause.Column{ | |||
| {Name: "user_id"}, | |||
| {Name: "channel_type"}, | |||
| {Name: "group"}, | |||
| }, | |||
| DoUpdates: clause.Assignments(map[string]interface{}{ | |||
| "channel_id": channelId, | |||
| "updated_at": now, | |||
| }), | |||
| }).Create(&binding).Error | |||
| } | |||
| func GetUserAssetChannelsByTypes(userId int, channelTypes []int, group string) ([]UserAssetChannel, error) { | |||
| if len(channelTypes) == 0 { | |||
| return nil, nil | |||
| } | |||
| var bindings []UserAssetChannel | |||
| err := DB.Where("user_id = ? AND channel_type IN ? AND "+commonGroupCol+" = ?", userId, channelTypes, group). | |||
| Order("updated_at DESC, id DESC"). | |||
| Find(&bindings).Error | |||
| return bindings, err | |||
| } | |||
| func DeleteUserAssetChannelsByTypesWithTx(tx *gorm.DB, userId int, channelTypes []int, group string) error { | |||
| if len(channelTypes) == 0 { | |||
| return nil | |||
| } | |||
| return tx.Session(&gorm.Session{SkipDefaultTransaction: true}).Where("user_id = ? AND channel_type IN ? AND "+commonGroupCol+" = ?", userId, channelTypes, group). | |||
| Delete(&UserAssetChannel{}).Error | |||
| } | |||
| func UnbindUserAssetChannel(userId int, channelType int, group string) error { | |||
| return DB.Where("user_id = ? AND channel_type = ? AND "+commonGroupCol+" = ?", userId, channelType, group). | |||
| Delete(&UserAssetChannel{}).Error | |||
| } | |||
| @@ -0,0 +1,221 @@ | |||
| package model | |||
| import ( | |||
| "sync" | |||
| "testing" | |||
| "github.com/glebarez/sqlite" | |||
| "github.com/stretchr/testify/assert" | |||
| "github.com/stretchr/testify/require" | |||
| "gorm.io/gorm" | |||
| "gorm.io/gorm/logger" | |||
| ) | |||
| func setupUserAssetChannelDB(t *testing.T) *gorm.DB { | |||
| t.Helper() | |||
| db, err := gorm.Open(sqlite.Open("file:user_asset_channels?mode=memory&cache=shared"), &gorm.Config{ | |||
| Logger: logger.Default.LogMode(logger.Silent), | |||
| }) | |||
| require.NoError(t, err) | |||
| sqlDB, err := db.DB() | |||
| require.NoError(t, err) | |||
| sqlDB.SetMaxOpenConns(1) | |||
| origDB := DB | |||
| origGroupCol := commonGroupCol | |||
| DB = db | |||
| commonGroupCol = "`group`" | |||
| require.NoError(t, db.AutoMigrate(&UserAssetChannel{})) | |||
| t.Cleanup(func() { | |||
| DB = origDB | |||
| commonGroupCol = origGroupCol | |||
| require.NoError(t, sqlDB.Close()) | |||
| }) | |||
| return db | |||
| } | |||
| func TestGetUserAssetChannel_NotFound(t *testing.T) { | |||
| setupUserAssetChannelDB(t) | |||
| binding, err := GetUserAssetChannel(1, 2, "default") | |||
| require.NoError(t, err) | |||
| assert.Nil(t, binding) | |||
| } | |||
| func TestBindUserAssetChannel_CreateAndUpdate(t *testing.T) { | |||
| db := setupUserAssetChannelDB(t) | |||
| require.NoError(t, BindUserAssetChannel(1, 2, "default", 100)) | |||
| binding, err := GetUserAssetChannel(1, 2, "default") | |||
| require.NoError(t, err) | |||
| require.NotNil(t, binding) | |||
| assert.Equal(t, 100, binding.ChannelId) | |||
| assert.NotZero(t, binding.CreatedAt) | |||
| assert.NotZero(t, binding.UpdatedAt) | |||
| require.NoError(t, BindUserAssetChannel(1, 2, "default", 200)) | |||
| binding, err = GetUserAssetChannel(1, 2, "default") | |||
| require.NoError(t, err) | |||
| require.NotNil(t, binding) | |||
| assert.Equal(t, 200, binding.ChannelId) | |||
| var count int64 | |||
| require.NoError(t, db.Model(&UserAssetChannel{}).Count(&count).Error) | |||
| assert.Equal(t, int64(1), count) | |||
| } | |||
| func TestGetUserAssetChannelsByTypesSortsLatestFirst(t *testing.T) { | |||
| db := setupUserAssetChannelDB(t) | |||
| require.NoError(t, BindUserAssetChannel(1, 2, "default", 100)) | |||
| require.NoError(t, BindUserAssetChannel(1, 3, "default", 200)) | |||
| require.NoError(t, BindUserAssetChannel(1, 4, "default", 300)) | |||
| require.NoError(t, db.Model(&UserAssetChannel{}). | |||
| Where("user_id = ? AND channel_type = ?", 1, 2). | |||
| Update("updated_at", int64(100)).Error) | |||
| require.NoError(t, db.Model(&UserAssetChannel{}). | |||
| Where("user_id = ? AND channel_type = ?", 1, 3). | |||
| Update("updated_at", int64(200)).Error) | |||
| bindings, err := GetUserAssetChannelsByTypes(1, []int{2, 3}, "default") | |||
| require.NoError(t, err) | |||
| require.Len(t, bindings, 2) | |||
| assert.Equal(t, 3, bindings[0].ChannelType) | |||
| assert.Equal(t, 200, bindings[0].ChannelId) | |||
| assert.Equal(t, 2, bindings[1].ChannelType) | |||
| } | |||
| func TestBindUserAssetChannelWithTxUpserts(t *testing.T) { | |||
| db := setupUserAssetChannelDB(t) | |||
| require.NoError(t, BindUserAssetChannel(1, 2, "default", 100)) | |||
| require.NoError(t, BindUserAssetChannel(1, 2, "default", 200)) | |||
| binding, err := GetUserAssetChannel(1, 2, "default") | |||
| require.NoError(t, err) | |||
| require.NotNil(t, binding) | |||
| assert.Equal(t, 200, binding.ChannelId) | |||
| var count int64 | |||
| require.NoError(t, db.Model(&UserAssetChannel{}).Count(&count).Error) | |||
| assert.Equal(t, int64(1), count) | |||
| } | |||
| func TestDeleteUserAssetChannelsByTypesWithTxDeletesOnlyRequestedRows(t *testing.T) { | |||
| db := setupUserAssetChannelDB(t) | |||
| require.NoError(t, BindUserAssetChannel(1, 2, "default", 100)) | |||
| require.NoError(t, BindUserAssetChannel(1, 3, "default", 200)) | |||
| require.NoError(t, BindUserAssetChannel(1, 4, "default", 300)) | |||
| require.NoError(t, BindUserAssetChannel(1, 2, "vip", 400)) | |||
| require.NoError(t, BindUserAssetChannel(2, 2, "default", 500)) | |||
| require.NoError(t, db.Transaction(func(tx *gorm.DB) error { | |||
| return DeleteUserAssetChannelsByTypesWithTx(tx, 1, []int{2, 3}, "default") | |||
| })) | |||
| binding, err := GetUserAssetChannel(1, 2, "default") | |||
| require.NoError(t, err) | |||
| assert.Nil(t, binding) | |||
| binding, err = GetUserAssetChannel(1, 3, "default") | |||
| require.NoError(t, err) | |||
| assert.Nil(t, binding) | |||
| binding, err = GetUserAssetChannel(1, 4, "default") | |||
| require.NoError(t, err) | |||
| require.NotNil(t, binding) | |||
| assert.Equal(t, 300, binding.ChannelId) | |||
| binding, err = GetUserAssetChannel(1, 2, "vip") | |||
| require.NoError(t, err) | |||
| require.NotNil(t, binding) | |||
| assert.Equal(t, 400, binding.ChannelId) | |||
| binding, err = GetUserAssetChannel(2, 2, "default") | |||
| require.NoError(t, err) | |||
| require.NotNil(t, binding) | |||
| assert.Equal(t, 500, binding.ChannelId) | |||
| } | |||
| func TestBindUserAssetChannel_IsolatesUserTypeGroup(t *testing.T) { | |||
| setupUserAssetChannelDB(t) | |||
| require.NoError(t, BindUserAssetChannel(1, 2, "default", 100)) | |||
| require.NoError(t, BindUserAssetChannel(2, 2, "default", 200)) | |||
| require.NoError(t, BindUserAssetChannel(1, 3, "default", 300)) | |||
| require.NoError(t, BindUserAssetChannel(1, 2, "vip", 400)) | |||
| cases := []struct { | |||
| userId int | |||
| channelType int | |||
| group string | |||
| channelId int | |||
| }{ | |||
| {1, 2, "default", 100}, | |||
| {2, 2, "default", 200}, | |||
| {1, 3, "default", 300}, | |||
| {1, 2, "vip", 400}, | |||
| } | |||
| for _, tc := range cases { | |||
| binding, err := GetUserAssetChannel(tc.userId, tc.channelType, tc.group) | |||
| require.NoError(t, err) | |||
| require.NotNil(t, binding) | |||
| assert.Equal(t, tc.channelId, binding.ChannelId) | |||
| } | |||
| } | |||
| func TestBindUserAssetChannel_ConcurrentUpsert(t *testing.T) { | |||
| db := setupUserAssetChannelDB(t) | |||
| const workers = 20 | |||
| var wg sync.WaitGroup | |||
| errCh := make(chan error, workers) | |||
| expected := make(map[int]bool, workers) | |||
| for i := 0; i < workers; i++ { | |||
| channelId := 1000 + i | |||
| expected[channelId] = true | |||
| wg.Add(1) | |||
| go func() { | |||
| defer wg.Done() | |||
| errCh <- BindUserAssetChannel(1, 2, "default", channelId) | |||
| }() | |||
| } | |||
| wg.Wait() | |||
| close(errCh) | |||
| for err := range errCh { | |||
| require.NoError(t, err) | |||
| } | |||
| var count int64 | |||
| require.NoError(t, db.Model(&UserAssetChannel{}).Count(&count).Error) | |||
| assert.Equal(t, int64(1), count) | |||
| binding, err := GetUserAssetChannel(1, 2, "default") | |||
| require.NoError(t, err) | |||
| require.NotNil(t, binding) | |||
| assert.True(t, expected[binding.ChannelId]) | |||
| } | |||
| func TestUnbindUserAssetChannel(t *testing.T) { | |||
| setupUserAssetChannelDB(t) | |||
| require.NoError(t, BindUserAssetChannel(1, 2, "default", 100)) | |||
| require.NoError(t, BindUserAssetChannel(1, 2, "vip", 200)) | |||
| require.NoError(t, UnbindUserAssetChannel(1, 2, "default")) | |||
| binding, err := GetUserAssetChannel(1, 2, "default") | |||
| require.NoError(t, err) | |||
| assert.Nil(t, binding) | |||
| binding, err = GetUserAssetChannel(1, 2, "vip") | |||
| require.NoError(t, err) | |||
| require.NotNil(t, binding) | |||
| assert.Equal(t, 200, binding.ChannelId) | |||
| } | |||
| @@ -778,14 +778,21 @@ func ClaudeStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon. | |||
| Usage: &dto.Usage{}, | |||
| } | |||
| var err *types.NewAPIError | |||
| var normalStreamStarted bool | |||
| helper.StreamScannerHandler(c, resp, info, func(data string) bool { | |||
| err = HandleStreamResponseData(c, info, claudeInfo, data) | |||
| if err != nil { | |||
| if normalStreamStarted || helper.StreamDataWritten(c) { | |||
| logger.LogError(c, "upstream claude stream error after normal stream started: "+err.Error()) | |||
| err = nil | |||
| } | |||
| return false | |||
| } | |||
| normalStreamStarted = true | |||
| return true | |||
| }) | |||
| if err != nil { | |||
| helper.ClearEventStreamHeadersIfNotWritten(c) | |||
| return nil, err | |||
| } | |||
| @@ -1,10 +1,18 @@ | |||
| package claude | |||
| import ( | |||
| "io" | |||
| "net/http" | |||
| "net/http/httptest" | |||
| "strings" | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/constant" | |||
| "github.com/QuantumNous/new-api/dto" | |||
| relaycommon "github.com/QuantumNous/new-api/relay/common" | |||
| "github.com/QuantumNous/new-api/types" | |||
| "github.com/gin-gonic/gin" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestFormatClaudeResponseInfo_MessageStart(t *testing.T) { | |||
| @@ -173,3 +181,65 @@ func TestFormatClaudeResponseInfo_ContentBlockDelta(t *testing.T) { | |||
| t.Errorf("ResponseText = %q, want %q", claudeInfo.ResponseText.String(), "hello") | |||
| } | |||
| } | |||
| func TestClaudeStreamHandlerReturnsErrorOnInitialErrorEvent(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| oldTimeout := constant.StreamingTimeout | |||
| constant.StreamingTimeout = 30 | |||
| t.Cleanup(func() { | |||
| constant.StreamingTimeout = oldTimeout | |||
| }) | |||
| w := httptest.NewRecorder() | |||
| c, _ := gin.CreateTestContext(w) | |||
| c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) | |||
| errorEvent := `{"type":"error","error":{"type":"overloaded_error","message":"Overloaded"}}` | |||
| resp := &http.Response{ | |||
| StatusCode: http.StatusOK, | |||
| Body: io.NopCloser(strings.NewReader("data: " + errorEvent + "\n\n")), | |||
| } | |||
| info := &relaycommon.RelayInfo{ | |||
| RelayFormat: types.RelayFormatClaude, | |||
| ChannelMeta: &relaycommon.ChannelMeta{ | |||
| UpstreamModelName: "claude-test", | |||
| }, | |||
| } | |||
| usage, err := ClaudeStreamHandler(c, resp, info) | |||
| require.Nil(t, usage) | |||
| require.Error(t, err) | |||
| require.Empty(t, w.Body.String()) | |||
| } | |||
| func TestClaudeStreamHandlerDoesNotRetryAfterNormalStreamData(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| oldTimeout := constant.StreamingTimeout | |||
| constant.StreamingTimeout = 30 | |||
| t.Cleanup(func() { | |||
| constant.StreamingTimeout = oldTimeout | |||
| }) | |||
| w := httptest.NewRecorder() | |||
| c, _ := gin.CreateTestContext(w) | |||
| c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) | |||
| normalEvent := `{"type":"message_start","message":{"id":"msg_test","model":"claude-test","usage":{"input_tokens":1,"output_tokens":1}}}` | |||
| errorEvent := `{"type":"error","error":{"type":"overloaded_error","message":"Overloaded"}}` | |||
| resp := &http.Response{ | |||
| StatusCode: http.StatusOK, | |||
| Body: io.NopCloser(strings.NewReader("data: " + normalEvent + "\n\ndata: " + errorEvent + "\n\n")), | |||
| } | |||
| info := &relaycommon.RelayInfo{ | |||
| RelayFormat: types.RelayFormatClaude, | |||
| ChannelMeta: &relaycommon.ChannelMeta{ | |||
| UpstreamModelName: "claude-test", | |||
| }, | |||
| } | |||
| usage, err := ClaudeStreamHandler(c, resp, info) | |||
| require.Nil(t, err) | |||
| require.NotNil(t, usage) | |||
| require.Contains(t, w.Body.String(), "message_start") | |||
| require.NotContains(t, w.Body.String(), "overloaded_error") | |||
| } | |||
| @@ -1285,6 +1285,8 @@ func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http | |||
| var usage = &dto.Usage{} | |||
| var imageCount int | |||
| responseText := strings.Builder{} | |||
| var streamErr *types.NewAPIError | |||
| var normalStreamStarted bool | |||
| helper.StreamScannerHandler(c, resp, info, func(data string) bool { | |||
| var geminiResponse dto.GeminiChatResponse | |||
| @@ -1296,8 +1298,21 @@ func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http | |||
| if len(geminiResponse.Candidates) == 0 && geminiResponse.PromptFeedback != nil && geminiResponse.PromptFeedback.BlockReason != nil { | |||
| common.SetContextKey(c, constant.ContextKeyAdminRejectReason, fmt.Sprintf("gemini_block_reason=%s", *geminiResponse.PromptFeedback.BlockReason)) | |||
| apiErr := types.NewOpenAIError( | |||
| errors.New("request blocked by Gemini API: "+*geminiResponse.PromptFeedback.BlockReason), | |||
| types.ErrorCodePromptBlocked, | |||
| http.StatusBadRequest, | |||
| ) | |||
| apiErr.UpstreamBody = service.TruncateBody(data) | |||
| if !normalStreamStarted && !helper.StreamDataWritten(c) { | |||
| streamErr = apiErr | |||
| } else { | |||
| logger.LogError(c, "upstream gemini stream blocked after normal stream started: "+apiErr.Error()) | |||
| } | |||
| return false | |||
| } | |||
| normalStreamStarted = true | |||
| // 统计图片数量 | |||
| for _, candidate := range geminiResponse.Candidates { | |||
| for _, part := range candidate.Content.Parts { | |||
| @@ -1318,6 +1333,10 @@ func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http | |||
| return callback(data, &geminiResponse) | |||
| }) | |||
| if streamErr != nil { | |||
| helper.ClearEventStreamHeadersIfNotWritten(c) | |||
| return nil, streamErr | |||
| } | |||
| if imageCount != 0 { | |||
| if usage.CompletionTokens == 0 { | |||
| @@ -282,6 +282,111 @@ func TestGeminiStreamHandlerUsesEstimatedPromptTokensWhenUsagePromptMissing(t *t | |||
| require.Equal(t, 110, usage.TotalTokens) | |||
| } | |||
| func TestGeminiStreamHandlerReturnsErrorOnPromptBlock(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| oldStreamingTimeout := constant.StreamingTimeout | |||
| constant.StreamingTimeout = 300 | |||
| t.Cleanup(func() { | |||
| constant.StreamingTimeout = oldStreamingTimeout | |||
| }) | |||
| w := httptest.NewRecorder() | |||
| c, _ := gin.CreateTestContext(w) | |||
| c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) | |||
| blockReason := "SAFETY" | |||
| payload := dto.GeminiChatResponse{ | |||
| PromptFeedback: &dto.GeminiChatPromptFeedback{ | |||
| BlockReason: &blockReason, | |||
| }, | |||
| } | |||
| body, err := common.Marshal(payload) | |||
| require.NoError(t, err) | |||
| resp := &http.Response{ | |||
| StatusCode: http.StatusOK, | |||
| Body: io.NopCloser(bytes.NewReader([]byte("data: " + string(body) + "\n\n"))), | |||
| } | |||
| info := &relaycommon.RelayInfo{ | |||
| RelayFormat: types.RelayFormatOpenAI, | |||
| ChannelMeta: &relaycommon.ChannelMeta{ | |||
| UpstreamModelName: "gemini-test", | |||
| }, | |||
| } | |||
| called := false | |||
| usage, newAPIError := geminiStreamHandler(c, info, resp, func(_ string, _ *dto.GeminiChatResponse) bool { | |||
| called = true | |||
| return true | |||
| }) | |||
| require.Nil(t, usage) | |||
| require.NotNil(t, newAPIError) | |||
| require.False(t, called) | |||
| require.Empty(t, w.Body.String()) | |||
| require.Empty(t, w.Header().Get("Content-Type")) | |||
| require.Contains(t, newAPIError.UpstreamBody, "SAFETY") | |||
| } | |||
| func TestGeminiStreamHandlerDoesNotRetryAfterNormalStreamData(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| oldStreamingTimeout := constant.StreamingTimeout | |||
| constant.StreamingTimeout = 300 | |||
| t.Cleanup(func() { | |||
| constant.StreamingTimeout = oldStreamingTimeout | |||
| }) | |||
| w := httptest.NewRecorder() | |||
| c, _ := gin.CreateTestContext(w) | |||
| c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) | |||
| normalChunk := dto.GeminiChatResponse{ | |||
| Candidates: []dto.GeminiChatCandidate{ | |||
| { | |||
| Content: dto.GeminiChatContent{ | |||
| Role: "model", | |||
| Parts: []dto.GeminiPart{ | |||
| {Text: "partial"}, | |||
| }, | |||
| }, | |||
| }, | |||
| }, | |||
| } | |||
| blockReason := "SAFETY" | |||
| blockChunk := dto.GeminiChatResponse{ | |||
| PromptFeedback: &dto.GeminiChatPromptFeedback{ | |||
| BlockReason: &blockReason, | |||
| }, | |||
| } | |||
| normalData, err := common.Marshal(normalChunk) | |||
| require.NoError(t, err) | |||
| blockData, err := common.Marshal(blockChunk) | |||
| require.NoError(t, err) | |||
| resp := &http.Response{ | |||
| StatusCode: http.StatusOK, | |||
| Body: io.NopCloser(bytes.NewReader([]byte("data: " + string(normalData) + "\n\ndata: " + string(blockData) + "\n\n"))), | |||
| } | |||
| info := &relaycommon.RelayInfo{ | |||
| RelayFormat: types.RelayFormatOpenAI, | |||
| ChannelMeta: &relaycommon.ChannelMeta{ | |||
| UpstreamModelName: "gemini-test", | |||
| }, | |||
| } | |||
| called := false | |||
| usage, newAPIError := geminiStreamHandler(c, info, resp, func(_ string, _ *dto.GeminiChatResponse) bool { | |||
| called = true | |||
| return true | |||
| }) | |||
| require.Nil(t, newAPIError) | |||
| require.NotNil(t, usage) | |||
| require.True(t, called) | |||
| require.Empty(t, w.Body.String()) | |||
| } | |||
| func TestGeminiTextGenerationHandlerUsesEstimatedPromptTokensWhenUsagePromptMissing(t *testing.T) { | |||
| t.Parallel() | |||
| @@ -2,6 +2,7 @@ package openai | |||
| import ( | |||
| "encoding/json" | |||
| "net/http" | |||
| "strings" | |||
| "github.com/QuantumNous/new-api/common" | |||
| @@ -19,6 +20,23 @@ import ( | |||
| ) | |||
| // 辅助函数 | |||
| func streamOpenAIErrorFromData(data string, statusCode int) *types.NewAPIError { | |||
| var errorResp dto.GeneralErrorResponse | |||
| if err := common.UnmarshalJsonStr(data, &errorResp); err != nil { | |||
| return nil | |||
| } | |||
| openaiErr := errorResp.TryToOpenAIError() | |||
| if openaiErr == nil { | |||
| return nil | |||
| } | |||
| if statusCode < 100 || statusCode > 599 || (statusCode >= 200 && statusCode < 300) { | |||
| statusCode = http.StatusInternalServerError | |||
| } | |||
| apiErr := types.WithOpenAIError(*openaiErr, statusCode) | |||
| apiErr.UpstreamBody = service.TruncateBody(data) | |||
| return apiErr | |||
| } | |||
| func HandleStreamFormat(c *gin.Context, info *relaycommon.RelayInfo, data string, forceFormat bool, thinkToContent bool) error { | |||
| info.SendResponseCount++ | |||
| @@ -121,12 +121,24 @@ func OaiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Re | |||
| var usage = &dto.Usage{} | |||
| var streamItems []string // store stream items | |||
| var lastStreamData string | |||
| var streamErr *types.NewAPIError | |||
| var normalStreamStarted bool | |||
| var secondLastStreamData string // 存储倒数第二个stream data,用于音频模型 | |||
| // 检查是否为音频模型 | |||
| isAudioModel := strings.Contains(strings.ToLower(model), "audio") | |||
| helper.StreamScannerHandler(c, resp, info, func(data string) bool { | |||
| if apiErr := streamOpenAIErrorFromData(data, resp.StatusCode); apiErr != nil { | |||
| if !normalStreamStarted && !helper.StreamDataWritten(c) { | |||
| streamErr = apiErr | |||
| } else { | |||
| logger.LogError(c, "upstream stream error after normal stream started: "+apiErr.Error()) | |||
| } | |||
| return false | |||
| } | |||
| normalStreamStarted = true | |||
| if lastStreamData != "" { | |||
| err := HandleStreamFormat(c, info, lastStreamData, info.ChannelSetting.ForceFormat, info.ChannelSetting.ThinkingToContent) | |||
| if err != nil { | |||
| @@ -146,6 +158,11 @@ func OaiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Re | |||
| }) | |||
| // 对音频模型,从倒数第二个stream data中提取usage信息 | |||
| if streamErr != nil { | |||
| helper.ClearEventStreamHeadersIfNotWritten(c) | |||
| return nil, streamErr | |||
| } | |||
| if isAudioModel && secondLastStreamData != "" { | |||
| var streamResp struct { | |||
| Usage *dto.Usage `json:"usage"` | |||
| @@ -85,11 +85,24 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp | |||
| var usage = &dto.Usage{} | |||
| var responseTextBuilder strings.Builder | |||
| var streamErr *types.NewAPIError | |||
| var normalStreamStarted bool | |||
| helper.StreamScannerHandler(c, resp, info, func(data string) bool { | |||
| var streamResponse dto.ResponsesStreamResponse | |||
| if err := common.UnmarshalJsonStr(data, &streamResponse); err == nil { | |||
| switch streamResponse.Type { | |||
| case "response.error", "response.failed", "error": | |||
| apiErr := handleResponsesStreamError(streamResponse, data) | |||
| if !normalStreamStarted && !helper.StreamDataWritten(c) { | |||
| streamErr = apiErr | |||
| } else if apiErr != nil { | |||
| logger.LogError(c, "upstream responses stream error after normal stream started: "+apiErr.Error()) | |||
| } | |||
| return false | |||
| } | |||
| normalStreamStarted = true | |||
| sendResponsesStreamData(c, streamResponse, data) | |||
| switch streamResponse.Type { | |||
| case "response.completed": | |||
| @@ -128,9 +141,6 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp | |||
| } | |||
| } | |||
| } | |||
| case "response.error", "response.failed", "error": | |||
| streamErr = handleResponsesStreamError(streamResponse, data) | |||
| return false | |||
| } | |||
| } else { | |||
| logger.LogError(c, "failed to unmarshal stream response: "+err.Error()) | |||
| @@ -139,6 +149,7 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp | |||
| }) | |||
| if streamErr != nil { | |||
| helper.ClearEventStreamHeadersIfNotWritten(c) | |||
| return nil, streamErr | |||
| } | |||
| @@ -8,8 +8,11 @@ import ( | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/constant" | |||
| "github.com/QuantumNous/new-api/dto" | |||
| relaycommon "github.com/QuantumNous/new-api/relay/common" | |||
| "github.com/QuantumNous/new-api/relay/helper" | |||
| "github.com/QuantumNous/new-api/types" | |||
| "github.com/gin-gonic/gin" | |||
| "github.com/stretchr/testify/assert" | |||
| "github.com/stretchr/testify/require" | |||
| @@ -86,6 +89,165 @@ func TestOaiResponsesHandlerAttachesUpstreamBodyOnInvalidJSON(t *testing.T) { | |||
| assert.Contains(t, err.UpstreamBody, `{"incomplete":`) | |||
| } | |||
| func TestOaiResponsesStreamHandlerDoesNotForwardErrorEvent(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| oldTimeout := constant.StreamingTimeout | |||
| constant.StreamingTimeout = 30 | |||
| t.Cleanup(func() { | |||
| constant.StreamingTimeout = oldTimeout | |||
| }) | |||
| w := httptest.NewRecorder() | |||
| c, _ := gin.CreateTestContext(w) | |||
| c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) | |||
| errorEvent := `{"type":"error","error":{"type":"tokens","code":"rate_limit_exceeded","message":"Request too large","param":null},"sequence_number":2}` | |||
| resp := &http.Response{ | |||
| StatusCode: http.StatusOK, | |||
| Body: io.NopCloser(strings.NewReader("data: " + errorEvent + "\n\n")), | |||
| } | |||
| usage, err := OaiResponsesStreamHandler(c, &relaycommon.RelayInfo{}, resp) | |||
| require.Nil(t, usage) | |||
| require.Error(t, err) | |||
| assert.Equal(t, http.StatusInternalServerError, err.StatusCode) | |||
| assert.Contains(t, err.UpstreamBody, "rate_limit_exceeded") | |||
| assert.NotContains(t, w.Body.String(), "rate_limit_exceeded") | |||
| assert.Empty(t, w.Body.String()) | |||
| assert.Empty(t, w.Header().Get("Content-Type")) | |||
| } | |||
| func TestOaiStreamHandlerDoesNotForwardInitialErrorEvent(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| oldTimeout := constant.StreamingTimeout | |||
| constant.StreamingTimeout = 30 | |||
| t.Cleanup(func() { | |||
| constant.StreamingTimeout = oldTimeout | |||
| }) | |||
| w := httptest.NewRecorder() | |||
| c, _ := gin.CreateTestContext(w) | |||
| c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) | |||
| errorEvent := `{"error":{"type":"tokens","code":"rate_limit_exceeded","message":"Request too large","param":null}}` | |||
| resp := &http.Response{ | |||
| StatusCode: http.StatusOK, | |||
| Body: io.NopCloser(strings.NewReader("data: " + errorEvent + "\n\n")), | |||
| } | |||
| info := &relaycommon.RelayInfo{ | |||
| RelayFormat: types.RelayFormatOpenAI, | |||
| ChannelMeta: &relaycommon.ChannelMeta{ | |||
| UpstreamModelName: "gpt-test", | |||
| }, | |||
| } | |||
| usage, err := OaiStreamHandler(c, info, resp) | |||
| require.Nil(t, usage) | |||
| require.Error(t, err) | |||
| assert.Contains(t, err.UpstreamBody, "rate_limit_exceeded") | |||
| assert.NotContains(t, w.Body.String(), "rate_limit_exceeded") | |||
| assert.Empty(t, w.Body.String()) | |||
| assert.Empty(t, w.Header().Get("Content-Type")) | |||
| } | |||
| func TestOaiStreamHandlerReturnsInitialErrorAfterPing(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| oldTimeout := constant.StreamingTimeout | |||
| constant.StreamingTimeout = 30 | |||
| t.Cleanup(func() { | |||
| constant.StreamingTimeout = oldTimeout | |||
| }) | |||
| w := httptest.NewRecorder() | |||
| c, _ := gin.CreateTestContext(w) | |||
| c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) | |||
| require.NoError(t, helper.PingData(c)) | |||
| errorEvent := `{"error":{"type":"tokens","code":"rate_limit_exceeded","message":"Request too large","param":null}}` | |||
| resp := &http.Response{ | |||
| StatusCode: http.StatusOK, | |||
| Body: io.NopCloser(strings.NewReader("data: " + errorEvent + "\n\n")), | |||
| } | |||
| info := &relaycommon.RelayInfo{ | |||
| RelayFormat: types.RelayFormatOpenAI, | |||
| ChannelMeta: &relaycommon.ChannelMeta{ | |||
| UpstreamModelName: "gpt-test", | |||
| }, | |||
| } | |||
| usage, err := OaiStreamHandler(c, info, resp) | |||
| require.Nil(t, usage) | |||
| require.Error(t, err) | |||
| assert.Equal(t, http.StatusInternalServerError, err.StatusCode) | |||
| assert.Contains(t, err.UpstreamBody, "rate_limit_exceeded") | |||
| assert.Contains(t, w.Body.String(), ": PING") | |||
| assert.NotContains(t, w.Body.String(), "rate_limit_exceeded") | |||
| } | |||
| func TestOaiStreamHandlerDoesNotRetryAfterNormalStreamData(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| oldTimeout := constant.StreamingTimeout | |||
| constant.StreamingTimeout = 30 | |||
| t.Cleanup(func() { | |||
| constant.StreamingTimeout = oldTimeout | |||
| }) | |||
| w := httptest.NewRecorder() | |||
| c, _ := gin.CreateTestContext(w) | |||
| c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) | |||
| normalEvent := `{"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,"model":"gpt-test","choices":[{"index":0,"delta":{"content":"hi"}}]}` | |||
| errorEvent := `{"error":{"type":"tokens","code":"rate_limit_exceeded","message":"Request too large","param":null}}` | |||
| resp := &http.Response{ | |||
| StatusCode: http.StatusOK, | |||
| Body: io.NopCloser(strings.NewReader("data: " + normalEvent + "\n\ndata: " + errorEvent + "\n\n")), | |||
| } | |||
| info := &relaycommon.RelayInfo{ | |||
| RelayFormat: types.RelayFormatOpenAI, | |||
| ChannelMeta: &relaycommon.ChannelMeta{ | |||
| UpstreamModelName: "gpt-test", | |||
| }, | |||
| } | |||
| usage, err := OaiStreamHandler(c, info, resp) | |||
| require.Nil(t, err) | |||
| require.NotNil(t, usage) | |||
| assert.Contains(t, w.Body.String(), "hi") | |||
| assert.NotContains(t, w.Body.String(), "rate_limit_exceeded") | |||
| } | |||
| func TestOaiResponsesStreamHandlerDoesNotRetryAfterNormalStreamData(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| oldTimeout := constant.StreamingTimeout | |||
| constant.StreamingTimeout = 30 | |||
| t.Cleanup(func() { | |||
| constant.StreamingTimeout = oldTimeout | |||
| }) | |||
| w := httptest.NewRecorder() | |||
| c, _ := gin.CreateTestContext(w) | |||
| c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) | |||
| normalEvent := `{"type":"response.completed","response":{"id":"resp_test","usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}}` | |||
| errorEvent := `{"type":"error","error":{"type":"tokens","code":"rate_limit_exceeded","message":"Request too large","param":null},"sequence_number":2}` | |||
| resp := &http.Response{ | |||
| StatusCode: http.StatusOK, | |||
| Body: io.NopCloser(strings.NewReader("data: " + normalEvent + "\n\ndata: " + errorEvent + "\n\n")), | |||
| } | |||
| info := &relaycommon.RelayInfo{ | |||
| ChannelMeta: &relaycommon.ChannelMeta{ | |||
| UpstreamModelName: "gpt-test", | |||
| }, | |||
| } | |||
| usage, err := OaiResponsesStreamHandler(c, info, resp) | |||
| require.Nil(t, err) | |||
| require.NotNil(t, usage) | |||
| assert.Contains(t, w.Body.String(), "response.completed") | |||
| assert.NotContains(t, w.Body.String(), "rate_limit_exceeded") | |||
| } | |||
| func TestOaiResponsesToChatHandlerAttachesUpstreamBodyOnInvalidJSON(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| w := httptest.NewRecorder() | |||
| @@ -5,6 +5,7 @@ import ( | |||
| "fmt" | |||
| "io" | |||
| "net/http" | |||
| "strconv" | |||
| "time" | |||
| "github.com/QuantumNous/new-api/common" | |||
| @@ -19,6 +20,7 @@ import ( | |||
| "github.com/gin-gonic/gin" | |||
| "github.com/pkg/errors" | |||
| "github.com/samber/lo" | |||
| ) | |||
| // ============================ | |||
| @@ -26,37 +28,37 @@ import ( | |||
| // ============================ | |||
| type ContentItem struct { | |||
| Type string `json:"type"` // "text", "image_url" or "video" | |||
| Text string `json:"text,omitempty"` // for text type | |||
| ImageURL *ImageURL `json:"image_url,omitempty"` // for image_url type | |||
| Video *VideoReference `json:"video,omitempty"` // for video (sample) type | |||
| Role string `json:"role,omitempty"` // reference_image / first_frame / last_frame | |||
| Type string `json:"type,omitempty"` | |||
| Text string `json:"text,omitempty"` | |||
| ImageURL *MediaURL `json:"image_url,omitempty"` | |||
| VideoURL *MediaURL `json:"video_url,omitempty"` | |||
| AudioURL *MediaURL `json:"audio_url,omitempty"` | |||
| Role string `json:"role,omitempty"` | |||
| } | |||
| type ImageURL struct { | |||
| URL string `json:"url"` | |||
| } | |||
| type VideoReference struct { | |||
| URL string `json:"url"` // Draft video URL | |||
| type MediaURL struct { | |||
| URL string `json:"url,omitempty"` | |||
| } | |||
| type requestPayload struct { | |||
| Model string `json:"model"` | |||
| Content []ContentItem `json:"content"` | |||
| Content []ContentItem `json:"content,omitempty"` | |||
| CallbackURL string `json:"callback_url,omitempty"` | |||
| ReturnLastFrame *dto.BoolValue `json:"return_last_frame,omitempty"` | |||
| ServiceTier string `json:"service_tier,omitempty"` | |||
| ExecutionExpiresAfter dto.IntValue `json:"execution_expires_after,omitempty"` | |||
| ExecutionExpiresAfter *dto.IntValue `json:"execution_expires_after,omitempty"` | |||
| GenerateAudio *dto.BoolValue `json:"generate_audio,omitempty"` | |||
| Draft *dto.BoolValue `json:"draft,omitempty"` | |||
| Resolution string `json:"resolution,omitempty"` | |||
| Ratio string `json:"ratio,omitempty"` | |||
| Duration dto.IntValue `json:"duration,omitempty"` | |||
| Frames dto.IntValue `json:"frames,omitempty"` | |||
| Seed dto.IntValue `json:"seed,omitempty"` | |||
| CameraFixed *dto.BoolValue `json:"camera_fixed,omitempty"` | |||
| Watermark *dto.BoolValue `json:"watermark,omitempty"` | |||
| Tools []struct { | |||
| Type string `json:"type,omitempty"` | |||
| } `json:"tools,omitempty"` | |||
| Resolution string `json:"resolution,omitempty"` | |||
| Ratio string `json:"ratio,omitempty"` | |||
| Duration *dto.IntValue `json:"duration,omitempty"` | |||
| Frames *dto.IntValue `json:"frames,omitempty"` | |||
| Seed *dto.IntValue `json:"seed,omitempty"` | |||
| CameraFixed *dto.BoolValue `json:"camera_fixed,omitempty"` | |||
| Watermark *dto.BoolValue `json:"watermark,omitempty"` | |||
| } | |||
| type responsePayload struct { | |||
| @@ -80,6 +82,10 @@ type responseTask struct { | |||
| CompletionTokens int `json:"completion_tokens"` | |||
| TotalTokens int `json:"total_tokens"` | |||
| } `json:"usage"` | |||
| Error struct { | |||
| Code string `json:"code"` | |||
| Message string `json:"message"` | |||
| } `json:"error"` | |||
| CreatedAt int64 `json:"created_at"` | |||
| UpdatedAt int64 `json:"updated_at"` | |||
| } | |||
| @@ -218,20 +224,12 @@ func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (* | |||
| 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{ | |||
| ImageURL: &MediaURL{ | |||
| URL: imgURL, | |||
| }, | |||
| }) | |||
| @@ -243,6 +241,16 @@ func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (* | |||
| return nil, errors.Wrap(err, "unmarshal metadata failed") | |||
| } | |||
| if sec, _ := strconv.Atoi(req.Seconds); sec > 0 { | |||
| r.Duration = lo.ToPtr(dto.IntValue(sec)) | |||
| } | |||
| r.Content = lo.Reject(r.Content, func(c ContentItem, _ int) bool { return c.Type == "text" }) | |||
| r.Content = append(r.Content, ContentItem{ | |||
| Type: "text", | |||
| Text: req.Prompt, | |||
| }) | |||
| return &r, nil | |||
| } | |||
| @@ -274,7 +282,10 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e | |||
| case "failed": | |||
| taskResult.Status = model.TaskStatusFailure | |||
| taskResult.Progress = "100%" | |||
| taskResult.Reason = "task failed" | |||
| taskResult.Reason = resTask.Error.Message | |||
| if taskResult.Reason == "" { | |||
| taskResult.Reason = "task failed" | |||
| } | |||
| default: | |||
| // Unknown status, treat as processing | |||
| taskResult.Status = model.TaskStatusInProgress | |||
| @@ -301,9 +312,17 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, erro | |||
| openAIVideo.Model = originTask.Properties.OriginModelName | |||
| if dResp.Status == "failed" { | |||
| message := dResp.Error.Message | |||
| if message == "" { | |||
| message = "task failed" | |||
| } | |||
| code := dResp.Error.Code | |||
| if code == "" { | |||
| code = "failed" | |||
| } | |||
| openAIVideo.Error = &dto.OpenAIVideoError{ | |||
| Message: "task failed", | |||
| Code: "failed", | |||
| Message: message, | |||
| Code: code, | |||
| } | |||
| } | |||
| @@ -0,0 +1,190 @@ | |||
| package doubao | |||
| import ( | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/dto" | |||
| "github.com/QuantumNous/new-api/model" | |||
| relaycommon "github.com/QuantumNous/new-api/relay/common" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestParseTaskResult_SuccessWithUsage(t *testing.T) { | |||
| adaptor := &TaskAdaptor{} | |||
| taskInfo, err := adaptor.ParseTaskResult([]byte(`{ | |||
| "status":"succeeded", | |||
| "content":{"video_url":"https://example.test/video.mp4"}, | |||
| "usage":{"completion_tokens":12,"total_tokens":3456} | |||
| }`)) | |||
| require.NoError(t, err) | |||
| require.Equal(t, string(model.TaskStatusSuccess), taskInfo.Status) | |||
| require.Equal(t, "https://example.test/video.mp4", taskInfo.Url) | |||
| require.Equal(t, 12, taskInfo.CompletionTokens) | |||
| require.Equal(t, 3456, taskInfo.TotalTokens) | |||
| } | |||
| func TestParseTaskResult_SuccessWithoutUsage(t *testing.T) { | |||
| adaptor := &TaskAdaptor{} | |||
| taskInfo, err := adaptor.ParseTaskResult([]byte(`{ | |||
| "status":"succeeded", | |||
| "content":{"video_url":"https://example.test/video.mp4"} | |||
| }`)) | |||
| require.NoError(t, err) | |||
| require.Equal(t, string(model.TaskStatusSuccess), taskInfo.Status) | |||
| require.Zero(t, taskInfo.TotalTokens) | |||
| } | |||
| func TestParseTaskResult_NonSuccessStatusesDoNotRequireUsage(t *testing.T) { | |||
| adaptor := &TaskAdaptor{} | |||
| cases := map[string]model.TaskStatus{ | |||
| "queued": model.TaskStatusQueued, | |||
| "processing": model.TaskStatusInProgress, | |||
| "failed": model.TaskStatusFailure, | |||
| } | |||
| for status, want := range cases { | |||
| taskInfo, err := adaptor.ParseTaskResult([]byte(`{"status":"` + status + `"}`)) | |||
| require.NoError(t, err) | |||
| require.Equal(t, string(want), taskInfo.Status) | |||
| require.Zero(t, taskInfo.TotalTokens) | |||
| } | |||
| } | |||
| func TestConvertToRequestPayload_MetadataSupportsVideoAudioAndTools(t *testing.T) { | |||
| adaptor := &TaskAdaptor{} | |||
| req := &relaycommon.TaskSubmitReq{ | |||
| Model: "doubao-seedance-2-0-260128", | |||
| Prompt: "make a reference video", | |||
| Metadata: map[string]interface{}{ | |||
| "content": []interface{}{ | |||
| map[string]interface{}{ | |||
| "type": "video_url", | |||
| "video_url": map[string]interface{}{ | |||
| "url": "https://example.test/input.mp4", | |||
| }, | |||
| "role": "reference_video", | |||
| }, | |||
| map[string]interface{}{ | |||
| "type": "audio_url", | |||
| "audio_url": map[string]interface{}{ | |||
| "url": "https://example.test/input.wav", | |||
| }, | |||
| "role": "reference_audio", | |||
| }, | |||
| }, | |||
| "tools": []interface{}{ | |||
| map[string]interface{}{"type": "web_search"}, | |||
| }, | |||
| }, | |||
| } | |||
| payload, err := adaptor.convertToRequestPayload(req) | |||
| require.NoError(t, err) | |||
| require.Len(t, payload.Content, 3) | |||
| require.Equal(t, "video_url", payload.Content[0].Type) | |||
| require.Equal(t, "https://example.test/input.mp4", payload.Content[0].VideoURL.URL) | |||
| require.Equal(t, "reference_video", payload.Content[0].Role) | |||
| require.Equal(t, "audio_url", payload.Content[1].Type) | |||
| require.Equal(t, "https://example.test/input.wav", payload.Content[1].AudioURL.URL) | |||
| require.Equal(t, "reference_audio", payload.Content[1].Role) | |||
| require.Len(t, payload.Tools, 1) | |||
| require.Equal(t, "web_search", payload.Tools[0].Type) | |||
| } | |||
| func TestConvertToRequestPayload_SecondsOverridesDuration(t *testing.T) { | |||
| adaptor := &TaskAdaptor{} | |||
| req := &relaycommon.TaskSubmitReq{ | |||
| Model: "doubao-seedance-1-5-pro-251215", | |||
| Prompt: "cat yawning", | |||
| Seconds: "8", | |||
| Metadata: map[string]interface{}{ | |||
| "duration": 5, | |||
| }, | |||
| } | |||
| payload, err := adaptor.convertToRequestPayload(req) | |||
| require.NoError(t, err) | |||
| require.NotNil(t, payload.Duration) | |||
| require.Equal(t, dto.IntValue(8), *payload.Duration) | |||
| } | |||
| func TestConvertToRequestPayload_PromptAppendedAfterMetadataAndReplacesMetadataText(t *testing.T) { | |||
| adaptor := &TaskAdaptor{} | |||
| req := &relaycommon.TaskSubmitReq{ | |||
| Model: "doubao-seedance-1-5-pro-251215", | |||
| Prompt: "current prompt", | |||
| Images: []string{"https://example.test/first.png"}, | |||
| Metadata: map[string]interface{}{ | |||
| "content": []interface{}{ | |||
| map[string]interface{}{ | |||
| "type": "text", | |||
| "text": "metadata prompt", | |||
| }, | |||
| map[string]interface{}{ | |||
| "type": "image_url", | |||
| "image_url": map[string]interface{}{ | |||
| "url": "https://example.test/second.png", | |||
| }, | |||
| "role": "last_frame", | |||
| }, | |||
| }, | |||
| }, | |||
| } | |||
| payload, err := adaptor.convertToRequestPayload(req) | |||
| require.NoError(t, err) | |||
| require.Len(t, payload.Content, 2) | |||
| require.Equal(t, "image_url", payload.Content[0].Type) | |||
| require.Equal(t, "https://example.test/second.png", payload.Content[0].ImageURL.URL) | |||
| require.Equal(t, "last_frame", payload.Content[0].Role) | |||
| require.Equal(t, "text", payload.Content[1].Type) | |||
| require.Equal(t, "current prompt", payload.Content[1].Text) | |||
| } | |||
| func TestParseTaskResult_FailedUsesUpstreamErrorMessage(t *testing.T) { | |||
| adaptor := &TaskAdaptor{} | |||
| taskInfo, err := adaptor.ParseTaskResult([]byte(`{ | |||
| "status":"failed", | |||
| "error":{"code":"InvalidParameter","message":"duration is invalid"} | |||
| }`)) | |||
| require.NoError(t, err) | |||
| require.Equal(t, string(model.TaskStatusFailure), taskInfo.Status) | |||
| require.Equal(t, "duration is invalid", taskInfo.Reason) | |||
| } | |||
| func TestConvertToOpenAIVideo_FailedUsesUpstreamError(t *testing.T) { | |||
| adaptor := &TaskAdaptor{} | |||
| task := &model.Task{ | |||
| TaskID: "task_public", | |||
| Status: model.TaskStatusFailure, | |||
| Progress: "100%", | |||
| CreatedAt: 100, | |||
| UpdatedAt: 200, | |||
| Properties: model.Properties{ | |||
| OriginModelName: "doubao-seedance-1-5-pro-251215", | |||
| }, | |||
| Data: []byte(`{ | |||
| "status":"failed", | |||
| "error":{"code":"InvalidParameter","message":"duration is invalid"} | |||
| }`), | |||
| } | |||
| data, err := adaptor.ConvertToOpenAIVideo(task) | |||
| require.NoError(t, err) | |||
| require.Contains(t, string(data), `"message":"duration is invalid"`) | |||
| require.Contains(t, string(data), `"code":"InvalidParameter"`) | |||
| } | |||
| func TestGetModelList_IncludesSeedance20Models(t *testing.T) { | |||
| adaptor := &TaskAdaptor{} | |||
| models := adaptor.GetModelList() | |||
| require.Contains(t, models, "doubao-seedance-2-0-260128") | |||
| require.Contains(t, models, "doubao-seedance-2-0-fast-260128") | |||
| } | |||
| @@ -5,6 +5,8 @@ var ModelList = []string{ | |||
| "doubao-seedance-1-0-lite-t2v", | |||
| "doubao-seedance-1-0-lite-i2v", | |||
| "doubao-seedance-1-5-pro-251215", | |||
| "doubao-seedance-2-0-260128", | |||
| "doubao-seedance-2-0-fast-260128", | |||
| } | |||
| var ChannelName = "doubao-video" | |||
| @@ -0,0 +1,316 @@ | |||
| package doubao_aiping | |||
| import ( | |||
| "bytes" | |||
| "fmt" | |||
| "io" | |||
| "net/http" | |||
| "strconv" | |||
| "strings" | |||
| "time" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/constant" | |||
| "github.com/QuantumNous/new-api/dto" | |||
| "github.com/QuantumNous/new-api/model" | |||
| "github.com/QuantumNous/new-api/relay/channel" | |||
| taskcommon "github.com/QuantumNous/new-api/relay/channel/task/taskcommon" | |||
| relaycommon "github.com/QuantumNous/new-api/relay/common" | |||
| "github.com/QuantumNous/new-api/service" | |||
| relaytypes "github.com/QuantumNous/new-api/types" | |||
| "github.com/gin-gonic/gin" | |||
| "github.com/pkg/errors" | |||
| "github.com/samber/lo" | |||
| ) | |||
| type ContentItem struct { | |||
| Type string `json:"type,omitempty"` | |||
| Text string `json:"text,omitempty"` | |||
| ImageURL *MediaURL `json:"image_url,omitempty"` | |||
| VideoURL *MediaURL `json:"video_url,omitempty"` | |||
| AudioURL *MediaURL `json:"audio_url,omitempty"` | |||
| DraftTask *DraftTask `json:"draft_task,omitempty"` | |||
| Role string `json:"role,omitempty"` | |||
| } | |||
| type MediaURL struct { | |||
| URL string `json:"url,omitempty"` | |||
| } | |||
| type DraftTask struct { | |||
| ID string `json:"id,omitempty"` | |||
| } | |||
| type requestPayload struct { | |||
| Model string `json:"model"` | |||
| Content []ContentItem `json:"content,omitempty"` | |||
| CallbackURL string `json:"callback_url,omitempty"` | |||
| ReturnLastFrame *dto.BoolValue `json:"return_last_frame,omitempty"` | |||
| ServiceTier string `json:"service_tier,omitempty"` | |||
| ExecutionExpiresAfter *dto.IntValue `json:"execution_expires_after,omitempty"` | |||
| GenerateAudio *dto.BoolValue `json:"generate_audio,omitempty"` | |||
| Draft *dto.BoolValue `json:"draft,omitempty"` | |||
| Tools []struct { | |||
| Type string `json:"type,omitempty"` | |||
| } `json:"tools,omitempty"` | |||
| Resolution string `json:"resolution,omitempty"` | |||
| Ratio string `json:"ratio,omitempty"` | |||
| Duration *dto.IntValue `json:"duration,omitempty"` | |||
| Frames *dto.IntValue `json:"frames,omitempty"` | |||
| Seed *dto.IntValue `json:"seed,omitempty"` | |||
| CameraFixed *dto.BoolValue `json:"camera_fixed,omitempty"` | |||
| Watermark *dto.BoolValue `json:"watermark,omitempty"` | |||
| SafetyIdentifier string `json:"safety_identifier,omitempty"` | |||
| Priority *dto.IntValue `json:"priority,omitempty"` | |||
| } | |||
| type responsePayload struct { | |||
| ID string `json:"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"` | |||
| ServiceTier string `json:"service_tier"` | |||
| Usage struct { | |||
| CompletionTokens int `json:"completion_tokens"` | |||
| TotalTokens int `json:"total_tokens"` | |||
| } `json:"usage"` | |||
| Error struct { | |||
| Code string `json:"code"` | |||
| Message string `json:"message"` | |||
| } `json:"error"` | |||
| CreatedAt int64 `json:"created_at"` | |||
| UpdatedAt int64 `json:"updated_at"` | |||
| } | |||
| type TaskAdaptor struct { | |||
| taskcommon.BaseBilling | |||
| ChannelType int | |||
| apiKey string | |||
| baseURL string | |||
| } | |||
| func (a *TaskAdaptor) Init(info *relaycommon.RelayInfo) { | |||
| a.ChannelType = info.ChannelType | |||
| a.baseURL = strings.TrimRight(info.ChannelBaseUrl, "/") | |||
| a.apiKey = info.ApiKey | |||
| } | |||
| func (a *TaskAdaptor) ValidateRequestAndSetAction(c *gin.Context, info *relaycommon.RelayInfo) (taskErr *dto.TaskError) { | |||
| if _, err := relaycommon.GetTaskRequest(c); err == nil { | |||
| info.Action = constant.TaskActionGenerate | |||
| return nil | |||
| } | |||
| return relaycommon.ValidateBasicTaskRequest(c, info, constant.TaskActionGenerate) | |||
| } | |||
| func (a *TaskAdaptor) BuildRequestURL(info *relaycommon.RelayInfo) (string, error) { | |||
| baseURL := a.baseURL | |||
| if strings.TrimSpace(baseURL) == "" && info != nil { | |||
| baseURL = strings.TrimRight(info.ChannelBaseUrl, "/") | |||
| } | |||
| return fmt.Sprintf("%s/api/v1/multimodal/sd/videos/contents/generations/tasks", baseURL), nil | |||
| } | |||
| 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 | |||
| } | |||
| func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayInfo) (io.Reader, error) { | |||
| req, err := relaycommon.GetTaskRequest(c) | |||
| if err != nil { | |||
| return nil, err | |||
| } | |||
| body, err := a.convertToRequestPayload(&req) | |||
| if err != nil { | |||
| return nil, errors.Wrap(err, "convert request payload failed") | |||
| } | |||
| if info.IsModelMapped { | |||
| body.Model = info.UpstreamModelName | |||
| } else { | |||
| info.UpstreamModelName = body.Model | |||
| } | |||
| data, err := common.Marshal(body) | |||
| if err != nil { | |||
| return nil, err | |||
| } | |||
| return bytes.NewReader(data), nil | |||
| } | |||
| func (a *TaskAdaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (*http.Response, error) { | |||
| return channel.DoTaskApiRequest(a, c, info, requestBody) | |||
| } | |||
| 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 { | |||
| return "", nil, service.TaskErrorWrapper(err, "read_response_body_failed", http.StatusInternalServerError) | |||
| } | |||
| _ = resp.Body.Close() | |||
| var dResp responsePayload | |||
| if err := common.Unmarshal(responseBody, &dResp); err != nil { | |||
| return "", nil, service.TaskErrorWrapper(errors.Wrapf(err, "body: %s", responseBody), "unmarshal_response_body_failed", http.StatusInternalServerError) | |||
| } | |||
| if strings.TrimSpace(dResp.ID) == "" { | |||
| return "", nil, service.TaskErrorWrapper(fmt.Errorf("task_id is empty"), "invalid_response", http.StatusInternalServerError) | |||
| } | |||
| clientPayload := map[string]any{} | |||
| if err := common.Unmarshal(responseBody, &clientPayload); err != nil { | |||
| return "", nil, service.TaskErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError) | |||
| } | |||
| clientPayload = relaytypes.CloneMapAny(clientPayload) | |||
| delete(clientPayload, "aiping_id") | |||
| clientPayload["id"] = info.PublicTaskID | |||
| if _, ok := clientPayload["created_at"]; !ok { | |||
| clientPayload["created_at"] = time.Now().Unix() | |||
| } | |||
| if _, ok := clientPayload["model"]; !ok { | |||
| clientPayload["model"] = info.OriginModelName | |||
| } | |||
| c.JSON(http.StatusOK, clientPayload) | |||
| return dResp.ID, responseBody, nil | |||
| } | |||
| func (a *TaskAdaptor) FetchTask(baseUrl, key string, body map[string]any, proxy string) (*http.Response, error) { | |||
| req, err := a.buildFetchRequest(baseUrl, key, body) | |||
| if err != nil { | |||
| return nil, err | |||
| } | |||
| client, err := service.GetHttpClientWithProxy(proxy) | |||
| if err != nil { | |||
| return nil, fmt.Errorf("new proxy http client failed: %w", err) | |||
| } | |||
| return client.Do(req) | |||
| } | |||
| func (a *TaskAdaptor) buildFetchRequest(baseUrl, key string, body map[string]any) (*http.Request, error) { | |||
| taskID, ok := body["task_id"].(string) | |||
| if !ok || strings.TrimSpace(taskID) == "" { | |||
| return nil, fmt.Errorf("invalid task_id") | |||
| } | |||
| uri := fmt.Sprintf("%s/api/v1/multimodal/sd/videos/contents/generations/tasks/%s", strings.TrimRight(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 req, nil | |||
| } | |||
| 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{}, | |||
| } | |||
| if req.HasImage() { | |||
| for _, imgURL := range req.Images { | |||
| r.Content = append(r.Content, ContentItem{ | |||
| Type: "image_url", | |||
| ImageURL: &MediaURL{URL: imgURL}, | |||
| }) | |||
| } | |||
| } | |||
| if err := taskcommon.UnmarshalMetadata(req.Metadata, &r); err != nil { | |||
| return nil, errors.Wrap(err, "unmarshal metadata failed") | |||
| } | |||
| if sec, _ := strconv.Atoi(req.Seconds); sec > 0 { | |||
| r.Duration = lo.ToPtr(dto.IntValue(sec)) | |||
| } | |||
| hasTextContent := lo.SomeBy(r.Content, func(c ContentItem) bool { | |||
| return c.Type == "text" && strings.TrimSpace(c.Text) != "" | |||
| }) | |||
| if !hasTextContent && strings.TrimSpace(req.Prompt) != "" { | |||
| r.Content = append(r.Content, ContentItem{Type: "text", Text: req.Prompt}) | |||
| } | |||
| return &r, nil | |||
| } | |||
| func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) { | |||
| resTask := responseTask{} | |||
| if err := common.Unmarshal(respBody, &resTask); err != nil { | |||
| return nil, errors.Wrap(err, "unmarshal task result failed") | |||
| } | |||
| taskResult := relaycommon.TaskInfo{Code: 0} | |||
| switch resTask.Status { | |||
| case "pending", "queued": | |||
| taskResult.Status = model.TaskStatusQueued | |||
| taskResult.Progress = "10%" | |||
| case "processing", "running": | |||
| taskResult.Status = model.TaskStatusInProgress | |||
| taskResult.Progress = "50%" | |||
| case "succeeded": | |||
| taskResult.Status = model.TaskStatusSuccess | |||
| taskResult.Progress = "100%" | |||
| taskResult.Url = resTask.Content.VideoURL | |||
| taskResult.CompletionTokens = resTask.Usage.CompletionTokens | |||
| taskResult.TotalTokens = resTask.Usage.TotalTokens | |||
| case "failed", "expired", "cancelled": | |||
| taskResult.Status = model.TaskStatusFailure | |||
| taskResult.Progress = "100%" | |||
| taskResult.Reason = resTask.Error.Message | |||
| if taskResult.Reason == "" { | |||
| taskResult.Reason = "task " + resTask.Status | |||
| } | |||
| default: | |||
| taskResult.Status = model.TaskStatusInProgress | |||
| taskResult.Progress = "30%" | |||
| } | |||
| return &taskResult, nil | |||
| } | |||
| func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, error) { | |||
| var dResp responseTask | |||
| if err := common.Unmarshal(originTask.Data, &dResp); err != nil { | |||
| return nil, errors.Wrap(err, "unmarshal aiping doubao task data failed") | |||
| } | |||
| openAIVideo := dto.NewOpenAIVideo() | |||
| openAIVideo.ID = originTask.TaskID | |||
| openAIVideo.TaskID = originTask.TaskID | |||
| openAIVideo.Status = originTask.Status.ToVideoStatus() | |||
| openAIVideo.SetProgressStr(originTask.Progress) | |||
| openAIVideo.SetMetadata("url", dResp.Content.VideoURL) | |||
| openAIVideo.CreatedAt = originTask.CreatedAt | |||
| openAIVideo.CompletedAt = originTask.UpdatedAt | |||
| openAIVideo.Model = originTask.Properties.OriginModelName | |||
| if dResp.Status == "failed" { | |||
| message := dResp.Error.Message | |||
| if message == "" { | |||
| message = "task failed" | |||
| } | |||
| code := dResp.Error.Code | |||
| if code == "" { | |||
| code = "failed" | |||
| } | |||
| openAIVideo.Error = &dto.OpenAIVideoError{Message: message, Code: code} | |||
| } | |||
| return common.Marshal(openAIVideo) | |||
| } | |||
| @@ -0,0 +1,273 @@ | |||
| package doubao_aiping | |||
| import ( | |||
| "io" | |||
| "net/http" | |||
| "net/http/httptest" | |||
| "strings" | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/constant" | |||
| "github.com/QuantumNous/new-api/model" | |||
| relaycommon "github.com/QuantumNous/new-api/relay/common" | |||
| "github.com/gin-gonic/gin" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestBuildRequestURLUsesAipingVideosEndpoint(t *testing.T) { | |||
| adaptor := &TaskAdaptor{} | |||
| adaptor.Init(&relaycommon.RelayInfo{ | |||
| ChannelMeta: &relaycommon.ChannelMeta{ | |||
| ChannelType: constant.ChannelTypeDoubaoVideoCompatibleAiping, | |||
| ChannelBaseUrl: "https://aiping.example.com", | |||
| ApiKey: "sk-test", | |||
| }, | |||
| }) | |||
| got, err := adaptor.BuildRequestURL(&relaycommon.RelayInfo{ | |||
| ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: "https://aiping.example.com"}, | |||
| }) | |||
| require.NoError(t, err) | |||
| require.Equal(t, "https://aiping.example.com/api/v1/multimodal/sd/videos/contents/generations/tasks", got) | |||
| } | |||
| func TestFetchTaskUsesAipingVideosEndpoint(t *testing.T) { | |||
| adaptor := &TaskAdaptor{} | |||
| req, err := adaptor.buildFetchRequest("https://aiping.example.com", "sk-test", map[string]any{"task_id": "cgt-123"}) | |||
| require.NoError(t, err) | |||
| require.Equal(t, http.MethodGet, req.Method) | |||
| require.Equal(t, "https://aiping.example.com/api/v1/multimodal/sd/videos/contents/generations/tasks/cgt-123", req.URL.String()) | |||
| require.Equal(t, "Bearer sk-test", req.Header.Get("Authorization")) | |||
| } | |||
| func TestBuildRequestBodyPassesNativeMetadataToUpstream(t *testing.T) { | |||
| adaptor := &TaskAdaptor{} | |||
| c := newTaskRequestContext(t, `{ | |||
| "model":"doubao-seedance-2-0-260128", | |||
| "prompt":"current prompt", | |||
| "metadata":{ | |||
| "content":[ | |||
| {"type":"video_url","video_url":{"url":"https://example.test/input.mp4"},"role":"reference_video"} | |||
| ], | |||
| "duration":5 | |||
| } | |||
| }`) | |||
| info := &relaycommon.RelayInfo{ | |||
| ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "doubao-seedance-2-0-260128"}, | |||
| } | |||
| body, err := adaptor.BuildRequestBody(c, info) | |||
| require.NoError(t, err) | |||
| data, err := io.ReadAll(body) | |||
| require.NoError(t, err) | |||
| require.Contains(t, string(data), `"video_url":{"url":"https://example.test/input.mp4"}`) | |||
| require.Contains(t, string(data), `"duration":5`) | |||
| require.Contains(t, string(data), `"text":"current prompt"`) | |||
| } | |||
| func TestBuildRequestBodyPreservesNativeInterleavedContentOrder(t *testing.T) { | |||
| adaptor := &TaskAdaptor{} | |||
| c := newTaskRequestContext(t, `{ | |||
| "model":"doubao-seedance-2-0-260128", | |||
| "metadata":{ | |||
| "content":[ | |||
| {"type":"text","text":"开场:海边日落"}, | |||
| {"type":"image_url","image_url":{"url":"asset://img1"},"role":"reference_image"}, | |||
| {"type":"text","text":"中段:车身环绕特写"}, | |||
| {"type":"video_url","video_url":{"url":"asset://vid1"},"role":"reference_video"}, | |||
| {"type":"text","text":"结尾:尾灯点亮"} | |||
| ], | |||
| "duration":5 | |||
| } | |||
| }`) | |||
| info := &relaycommon.RelayInfo{ | |||
| ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "doubao-seedance-2-0-260128"}, | |||
| } | |||
| body, err := adaptor.BuildRequestBody(c, info) | |||
| require.NoError(t, err) | |||
| data, err := io.ReadAll(body) | |||
| require.NoError(t, err) | |||
| var payload requestPayload | |||
| require.NoError(t, common.Unmarshal(data, &payload)) | |||
| require.Len(t, payload.Content, 5) | |||
| require.Equal(t, "text", payload.Content[0].Type) | |||
| require.Equal(t, "开场:海边日落", payload.Content[0].Text) | |||
| require.Equal(t, "image_url", payload.Content[1].Type) | |||
| require.Equal(t, "asset://img1", payload.Content[1].ImageURL.URL) | |||
| require.Equal(t, "text", payload.Content[2].Type) | |||
| require.Equal(t, "中段:车身环绕特写", payload.Content[2].Text) | |||
| require.Equal(t, "video_url", payload.Content[3].Type) | |||
| require.Equal(t, "asset://vid1", payload.Content[3].VideoURL.URL) | |||
| require.Equal(t, "text", payload.Content[4].Type) | |||
| require.Equal(t, "结尾:尾灯点亮", payload.Content[4].Text) | |||
| } | |||
| func TestConvertToRequestPayloadPreservesNativeContentTextOrder(t *testing.T) { | |||
| adaptor := &TaskAdaptor{} | |||
| req := &relaycommon.TaskSubmitReq{ | |||
| Model: "doubao-seedance-2-0-260128", | |||
| Prompt: "first prompt\nsecond prompt\nthird prompt", | |||
| Metadata: map[string]interface{}{ | |||
| "content": []interface{}{ | |||
| map[string]interface{}{"type": "text", "text": "first prompt"}, | |||
| map[string]interface{}{"type": "image_url", "image_url": map[string]interface{}{"url": "asset://img1"}, "role": "reference_image"}, | |||
| map[string]interface{}{"type": "text", "text": "second prompt"}, | |||
| map[string]interface{}{"type": "video_url", "video_url": map[string]interface{}{"url": "asset://vid1"}, "role": "reference_video"}, | |||
| map[string]interface{}{"type": "text", "text": "third prompt"}, | |||
| }, | |||
| }, | |||
| } | |||
| body, err := adaptor.convertToRequestPayload(req) | |||
| require.NoError(t, err) | |||
| require.Len(t, body.Content, 5) | |||
| require.Equal(t, "text", body.Content[0].Type) | |||
| require.Equal(t, "first prompt", body.Content[0].Text) | |||
| require.Equal(t, "image_url", body.Content[1].Type) | |||
| require.Equal(t, "text", body.Content[2].Type) | |||
| require.Equal(t, "second prompt", body.Content[2].Text) | |||
| require.Equal(t, "video_url", body.Content[3].Type) | |||
| require.Equal(t, "text", body.Content[4].Type) | |||
| require.Equal(t, "third prompt", body.Content[4].Text) | |||
| } | |||
| func TestConvertToRequestPayloadPassesArkDocumentFields(t *testing.T) { | |||
| adaptor := &TaskAdaptor{} | |||
| req := &relaycommon.TaskSubmitReq{ | |||
| Model: "doubao-seedance-2-0-260128", | |||
| Prompt: "documented prompt", | |||
| Metadata: map[string]interface{}{ | |||
| "content": []interface{}{ | |||
| map[string]interface{}{"type": "text", "text": "documented prompt"}, | |||
| map[string]interface{}{"type": "image_url", "image_url": map[string]interface{}{"url": "asset://image-1"}, "role": "reference_image"}, | |||
| map[string]interface{}{"type": "video_url", "video_url": map[string]interface{}{"url": "asset://video-1"}, "role": "reference_video"}, | |||
| map[string]interface{}{"type": "audio_url", "audio_url": map[string]interface{}{"url": "asset://audio-1"}, "role": "reference_audio"}, | |||
| map[string]interface{}{"type": "draft_task", "draft_task": map[string]interface{}{"id": "cgt-draft"}}, | |||
| }, | |||
| "callback_url": "https://example.test/callback", | |||
| "return_last_frame": true, | |||
| "service_tier": "default", | |||
| "execution_expires_after": float64(3600), | |||
| "generate_audio": false, | |||
| "draft": true, | |||
| "tools": []interface{}{map[string]interface{}{"type": "web_search"}}, | |||
| "safety_identifier": "user-hash-1", | |||
| "priority": float64(5), | |||
| "resolution": "480p", | |||
| "ratio": "1:1", | |||
| "duration": float64(5), | |||
| "frames": float64(29), | |||
| "seed": float64(11), | |||
| "camera_fixed": false, | |||
| "watermark": true, | |||
| }, | |||
| } | |||
| body, err := adaptor.convertToRequestPayload(req) | |||
| require.NoError(t, err) | |||
| data, err := common.Marshal(body) | |||
| require.NoError(t, err) | |||
| jsonBody := string(data) | |||
| require.Contains(t, jsonBody, `"image_url":{"url":"asset://image-1"}`) | |||
| require.Contains(t, jsonBody, `"video_url":{"url":"asset://video-1"}`) | |||
| require.Contains(t, jsonBody, `"audio_url":{"url":"asset://audio-1"}`) | |||
| require.Contains(t, jsonBody, `"draft_task":{"id":"cgt-draft"}`) | |||
| require.Contains(t, jsonBody, `"callback_url":"https://example.test/callback"`) | |||
| require.Contains(t, jsonBody, `"return_last_frame":true`) | |||
| require.Contains(t, jsonBody, `"service_tier":"default"`) | |||
| require.Contains(t, jsonBody, `"execution_expires_after":3600`) | |||
| require.Contains(t, jsonBody, `"generate_audio":false`) | |||
| require.Contains(t, jsonBody, `"draft":true`) | |||
| require.Contains(t, jsonBody, `"tools":[{"type":"web_search"}]`) | |||
| require.Contains(t, jsonBody, `"safety_identifier":"user-hash-1"`) | |||
| require.Contains(t, jsonBody, `"priority":5`) | |||
| require.Contains(t, jsonBody, `"resolution":"480p"`) | |||
| require.Contains(t, jsonBody, `"ratio":"1:1"`) | |||
| require.Contains(t, jsonBody, `"duration":5`) | |||
| require.Contains(t, jsonBody, `"frames":29`) | |||
| require.Contains(t, jsonBody, `"seed":11`) | |||
| require.Contains(t, jsonBody, `"camera_fixed":false`) | |||
| require.Contains(t, jsonBody, `"watermark":true`) | |||
| } | |||
| func TestDoResponseRewritesIDToPublicTaskID(t *testing.T) { | |||
| adaptor := &TaskAdaptor{} | |||
| w := httptest.NewRecorder() | |||
| c, _ := gin.CreateTestContext(w) | |||
| info := &relaycommon.RelayInfo{ | |||
| OriginModelName: "doubao-seedance-2-0-260128", | |||
| TaskRelayInfo: &relaycommon.TaskRelayInfo{PublicTaskID: "task_public"}, | |||
| } | |||
| resp := &http.Response{ | |||
| StatusCode: http.StatusOK, | |||
| Body: io.NopCloser(strings.NewReader(`{ | |||
| "id":"cgt-upstream", | |||
| "aiping_id":"internal-aiping-id", | |||
| "status":"queued", | |||
| "created_at":1781496040, | |||
| "model":"doubao-seedance-2-0-260128" | |||
| }`)), | |||
| } | |||
| upstreamID, data, taskErr := adaptor.DoResponse(c, resp, info) | |||
| require.Nil(t, taskErr) | |||
| require.Equal(t, "cgt-upstream", upstreamID) | |||
| require.Contains(t, string(data), `"id":"cgt-upstream"`) | |||
| require.Equal(t, http.StatusOK, w.Code) | |||
| require.Contains(t, w.Body.String(), `"id":"task_public"`) | |||
| require.NotContains(t, w.Body.String(), "aiping_id") | |||
| require.Contains(t, w.Body.String(), `"status":"queued"`) | |||
| require.Contains(t, w.Body.String(), `"model":"doubao-seedance-2-0-260128"`) | |||
| } | |||
| func TestParseTaskResultTreatsArkTerminalStatusesAsFinished(t *testing.T) { | |||
| adaptor := &TaskAdaptor{} | |||
| cases := []struct { | |||
| name string | |||
| body string | |||
| wantReason string | |||
| }{ | |||
| { | |||
| name: "expired", | |||
| body: `{"id":"cgt-expired","status":"expired","error":{"code":"Expired","message":"task expired"}}`, | |||
| wantReason: "task expired", | |||
| }, | |||
| { | |||
| name: "cancelled", | |||
| body: `{"id":"cgt-cancelled","status":"cancelled","error":{"code":"Cancelled","message":"task cancelled"}}`, | |||
| wantReason: "task cancelled", | |||
| }, | |||
| } | |||
| for _, tc := range cases { | |||
| t.Run(tc.name, func(t *testing.T) { | |||
| got, err := adaptor.ParseTaskResult([]byte(tc.body)) | |||
| require.NoError(t, err) | |||
| require.Equal(t, model.TaskStatusFailure, got.Status) | |||
| require.Equal(t, "100%", got.Progress) | |||
| require.Equal(t, tc.wantReason, got.Reason) | |||
| }) | |||
| } | |||
| } | |||
| func newTaskRequestContext(t *testing.T, body string) *gin.Context { | |||
| t.Helper() | |||
| w := httptest.NewRecorder() | |||
| c, _ := gin.CreateTestContext(w) | |||
| c.Request = httptest.NewRequest(http.MethodPost, "/api/v3/contents/generations/tasks", strings.NewReader(body)) | |||
| c.Request.Header.Set("Content-Type", "application/json") | |||
| var req relaycommon.TaskSubmitReq | |||
| require.NoError(t, common.Unmarshal([]byte(body), &req)) | |||
| relaycommon.StoreTaskRequest(c, &relaycommon.RelayInfo{}, constant.TaskActionGenerate, req) | |||
| return c | |||
| } | |||
| @@ -0,0 +1,8 @@ | |||
| package doubao_aiping | |||
| var ModelList = []string{ | |||
| "doubao-seedance-2-0-260128", | |||
| "doubao-seedance-2-0-fast-260128", | |||
| } | |||
| var ChannelName = "DoubaoVideoCompatibleAiping" | |||
| @@ -0,0 +1,308 @@ | |||
| package doubao_tianyiyun | |||
| import ( | |||
| "bytes" | |||
| "fmt" | |||
| "io" | |||
| "net/http" | |||
| "strconv" | |||
| "strings" | |||
| "time" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/constant" | |||
| "github.com/QuantumNous/new-api/dto" | |||
| "github.com/QuantumNous/new-api/model" | |||
| "github.com/QuantumNous/new-api/relay/channel" | |||
| taskcommon "github.com/QuantumNous/new-api/relay/channel/task/taskcommon" | |||
| relaycommon "github.com/QuantumNous/new-api/relay/common" | |||
| "github.com/QuantumNous/new-api/service" | |||
| relaytypes "github.com/QuantumNous/new-api/types" | |||
| "github.com/gin-gonic/gin" | |||
| "github.com/pkg/errors" | |||
| "github.com/samber/lo" | |||
| ) | |||
| type ContentItem struct { | |||
| Type string `json:"type,omitempty"` | |||
| Text string `json:"text,omitempty"` | |||
| ImageURL *MediaURL `json:"image_url,omitempty"` | |||
| VideoURL *MediaURL `json:"video_url,omitempty"` | |||
| AudioURL *MediaURL `json:"audio_url,omitempty"` | |||
| DraftTask *DraftTask `json:"draft_task,omitempty"` | |||
| Role string `json:"role,omitempty"` | |||
| } | |||
| type MediaURL struct { | |||
| URL string `json:"url,omitempty"` | |||
| } | |||
| type DraftTask struct { | |||
| ID string `json:"id,omitempty"` | |||
| } | |||
| type requestPayload struct { | |||
| Model string `json:"model"` | |||
| Content []ContentItem `json:"content,omitempty"` | |||
| CallbackURL string `json:"callback_url,omitempty"` | |||
| ReturnLastFrame *dto.BoolValue `json:"return_last_frame,omitempty"` | |||
| ServiceTier string `json:"service_tier,omitempty"` | |||
| ExecutionExpiresAfter *dto.IntValue `json:"execution_expires_after,omitempty"` | |||
| GenerateAudio *dto.BoolValue `json:"generate_audio,omitempty"` | |||
| Draft *dto.BoolValue `json:"draft,omitempty"` | |||
| Tools []struct { | |||
| Type string `json:"type,omitempty"` | |||
| } `json:"tools,omitempty"` | |||
| Resolution string `json:"resolution,omitempty"` | |||
| Ratio string `json:"ratio,omitempty"` | |||
| Duration *dto.IntValue `json:"duration,omitempty"` | |||
| Frames *dto.IntValue `json:"frames,omitempty"` | |||
| Seed *dto.IntValue `json:"seed,omitempty"` | |||
| CameraFixed *dto.BoolValue `json:"camera_fixed,omitempty"` | |||
| Watermark *dto.BoolValue `json:"watermark,omitempty"` | |||
| SafetyIdentifier string `json:"safety_identifier,omitempty"` | |||
| Priority *dto.IntValue `json:"priority,omitempty"` | |||
| } | |||
| type responsePayload struct { | |||
| ID string `json:"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"` | |||
| Usage struct { | |||
| CompletionTokens int `json:"completion_tokens"` | |||
| TotalTokens int `json:"total_tokens"` | |||
| } `json:"usage"` | |||
| Error struct { | |||
| Code string `json:"code"` | |||
| Message string `json:"message"` | |||
| } `json:"error"` | |||
| CreatedAt int64 `json:"created_at"` | |||
| UpdatedAt int64 `json:"updated_at"` | |||
| } | |||
| type TaskAdaptor struct { | |||
| taskcommon.BaseBilling | |||
| ChannelType int | |||
| apiKey string | |||
| baseURL string | |||
| } | |||
| func (a *TaskAdaptor) Init(info *relaycommon.RelayInfo) { | |||
| a.ChannelType = info.ChannelType | |||
| a.baseURL = strings.TrimRight(info.ChannelBaseUrl, "/") | |||
| a.apiKey = info.ApiKey | |||
| } | |||
| func (a *TaskAdaptor) ValidateRequestAndSetAction(c *gin.Context, info *relaycommon.RelayInfo) (taskErr *dto.TaskError) { | |||
| if _, err := relaycommon.GetTaskRequest(c); err == nil { | |||
| info.Action = constant.TaskActionGenerate | |||
| return nil | |||
| } | |||
| return relaycommon.ValidateBasicTaskRequest(c, info, constant.TaskActionGenerate) | |||
| } | |||
| func (a *TaskAdaptor) BuildRequestURL(info *relaycommon.RelayInfo) (string, error) { | |||
| baseURL := a.baseURL | |||
| if strings.TrimSpace(baseURL) == "" && info != nil { | |||
| baseURL = strings.TrimRight(info.ChannelBaseUrl, "/") | |||
| } | |||
| return fmt.Sprintf("%s/v1/contents/generations/tasks", baseURL), nil | |||
| } | |||
| 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 | |||
| } | |||
| func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayInfo) (io.Reader, error) { | |||
| req, err := relaycommon.GetTaskRequest(c) | |||
| if err != nil { | |||
| return nil, err | |||
| } | |||
| body, err := a.convertToRequestPayload(&req) | |||
| if err != nil { | |||
| return nil, errors.Wrap(err, "convert request payload failed") | |||
| } | |||
| if info.IsModelMapped { | |||
| body.Model = info.UpstreamModelName | |||
| } else { | |||
| info.UpstreamModelName = body.Model | |||
| } | |||
| data, err := common.Marshal(body) | |||
| if err != nil { | |||
| return nil, err | |||
| } | |||
| return bytes.NewReader(data), nil | |||
| } | |||
| func (a *TaskAdaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (*http.Response, error) { | |||
| return channel.DoTaskApiRequest(a, c, info, requestBody) | |||
| } | |||
| func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (string, []byte, *dto.TaskError) { | |||
| responseBody, err := io.ReadAll(resp.Body) | |||
| if err != nil { | |||
| return "", nil, service.TaskErrorWrapper(err, "read_response_body_failed", http.StatusInternalServerError) | |||
| } | |||
| _ = resp.Body.Close() | |||
| var dResp responsePayload | |||
| if err := common.Unmarshal(responseBody, &dResp); err != nil { | |||
| return "", nil, service.TaskErrorWrapper(errors.Wrapf(err, "body: %s", responseBody), "unmarshal_response_body_failed", http.StatusInternalServerError) | |||
| } | |||
| if strings.TrimSpace(dResp.ID) == "" { | |||
| return "", nil, service.TaskErrorWrapper(fmt.Errorf("task_id is empty"), "invalid_response", http.StatusInternalServerError) | |||
| } | |||
| clientPayload := map[string]any{} | |||
| if err := common.Unmarshal(responseBody, &clientPayload); err != nil { | |||
| return "", nil, service.TaskErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError) | |||
| } | |||
| clientPayload = relaytypes.CloneMapAny(clientPayload) | |||
| clientPayload["id"] = info.PublicTaskID | |||
| if _, ok := clientPayload["created_at"]; !ok { | |||
| clientPayload["created_at"] = time.Now().Unix() | |||
| } | |||
| if _, ok := clientPayload["model"]; !ok { | |||
| clientPayload["model"] = info.OriginModelName | |||
| } | |||
| c.JSON(http.StatusOK, clientPayload) | |||
| return dResp.ID, responseBody, nil | |||
| } | |||
| func (a *TaskAdaptor) FetchTask(baseUrl, key string, body map[string]any, proxy string) (*http.Response, error) { | |||
| req, err := a.buildFetchRequest(baseUrl, key, body) | |||
| if err != nil { | |||
| return nil, err | |||
| } | |||
| client, err := service.GetHttpClientWithProxy(proxy) | |||
| if err != nil { | |||
| return nil, fmt.Errorf("new proxy http client failed: %w", err) | |||
| } | |||
| return client.Do(req) | |||
| } | |||
| func (a *TaskAdaptor) buildFetchRequest(baseUrl, key string, body map[string]any) (*http.Request, error) { | |||
| taskID, ok := body["task_id"].(string) | |||
| if !ok || strings.TrimSpace(taskID) == "" { | |||
| return nil, fmt.Errorf("invalid task_id") | |||
| } | |||
| uri := fmt.Sprintf("%s/v1/contents/generations/tasks/%s", strings.TrimRight(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 req, nil | |||
| } | |||
| 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{}, | |||
| } | |||
| if req.HasImage() { | |||
| for _, imgURL := range req.Images { | |||
| r.Content = append(r.Content, ContentItem{ | |||
| Type: "image_url", | |||
| ImageURL: &MediaURL{URL: imgURL}, | |||
| }) | |||
| } | |||
| } | |||
| if err := taskcommon.UnmarshalMetadata(req.Metadata, &r); err != nil { | |||
| return nil, errors.Wrap(err, "unmarshal metadata failed") | |||
| } | |||
| if sec, _ := strconv.Atoi(req.Seconds); sec > 0 && r.Duration == nil { | |||
| r.Duration = lo.ToPtr(dto.IntValue(sec)) | |||
| } | |||
| hasTextContent := lo.SomeBy(r.Content, func(c ContentItem) bool { | |||
| return c.Type == "text" && strings.TrimSpace(c.Text) != "" | |||
| }) | |||
| if !hasTextContent && strings.TrimSpace(req.Prompt) != "" { | |||
| r.Content = append(r.Content, ContentItem{Type: "text", Text: req.Prompt}) | |||
| } | |||
| return &r, nil | |||
| } | |||
| func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) { | |||
| resTask := responseTask{} | |||
| if err := common.Unmarshal(respBody, &resTask); err != nil { | |||
| return nil, errors.Wrap(err, "unmarshal task result failed") | |||
| } | |||
| taskResult := relaycommon.TaskInfo{Code: 0} | |||
| switch resTask.Status { | |||
| case "pending", "queued": | |||
| taskResult.Status = model.TaskStatusQueued | |||
| taskResult.Progress = "10%" | |||
| case "processing", "running": | |||
| taskResult.Status = model.TaskStatusInProgress | |||
| taskResult.Progress = "50%" | |||
| case "succeeded", "success": | |||
| taskResult.Status = model.TaskStatusSuccess | |||
| taskResult.Progress = "100%" | |||
| taskResult.Url = resTask.Content.VideoURL | |||
| taskResult.CompletionTokens = resTask.Usage.CompletionTokens | |||
| taskResult.TotalTokens = resTask.Usage.TotalTokens | |||
| case "failed", "expired", "cancelled": | |||
| taskResult.Status = model.TaskStatusFailure | |||
| taskResult.Progress = "100%" | |||
| taskResult.Reason = resTask.Error.Message | |||
| if taskResult.Reason == "" { | |||
| taskResult.Reason = "task " + resTask.Status | |||
| } | |||
| default: | |||
| taskResult.Status = model.TaskStatusInProgress | |||
| taskResult.Progress = "30%" | |||
| } | |||
| return &taskResult, nil | |||
| } | |||
| func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, error) { | |||
| var dResp responseTask | |||
| if err := common.Unmarshal(originTask.Data, &dResp); err != nil { | |||
| return nil, errors.Wrap(err, "unmarshal tianyiyun doubao task data failed") | |||
| } | |||
| openAIVideo := dto.NewOpenAIVideo() | |||
| openAIVideo.ID = originTask.TaskID | |||
| openAIVideo.TaskID = originTask.TaskID | |||
| openAIVideo.Status = originTask.Status.ToVideoStatus() | |||
| openAIVideo.SetProgressStr(originTask.Progress) | |||
| openAIVideo.SetMetadata("url", dResp.Content.VideoURL) | |||
| openAIVideo.CreatedAt = originTask.CreatedAt | |||
| openAIVideo.CompletedAt = originTask.UpdatedAt | |||
| openAIVideo.Model = originTask.Properties.OriginModelName | |||
| if dResp.Status == "failed" { | |||
| message := dResp.Error.Message | |||
| if message == "" { | |||
| message = "task failed" | |||
| } | |||
| code := dResp.Error.Code | |||
| if code == "" { | |||
| code = "failed" | |||
| } | |||
| openAIVideo.Error = &dto.OpenAIVideoError{Message: message, Code: code} | |||
| } | |||
| return common.Marshal(openAIVideo) | |||
| } | |||
| @@ -0,0 +1,206 @@ | |||
| package doubao_tianyiyun | |||
| import ( | |||
| "io" | |||
| "net/http" | |||
| "net/http/httptest" | |||
| "strings" | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/constant" | |||
| "github.com/QuantumNous/new-api/model" | |||
| relaycommon "github.com/QuantumNous/new-api/relay/common" | |||
| "github.com/gin-gonic/gin" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestBuildRequestURLUsesTianyiYunEndpoint(t *testing.T) { | |||
| adaptor := &TaskAdaptor{} | |||
| adaptor.Init(&relaycommon.RelayInfo{ | |||
| ChannelMeta: &relaycommon.ChannelMeta{ | |||
| ChannelType: constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, | |||
| ChannelBaseUrl: "https://ai.ctaigw.cn", | |||
| ApiKey: "sk-test", | |||
| }, | |||
| }) | |||
| got, err := adaptor.BuildRequestURL(&relaycommon.RelayInfo{ | |||
| ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: "https://ai.ctaigw.cn"}, | |||
| }) | |||
| require.NoError(t, err) | |||
| require.Equal(t, "https://ai.ctaigw.cn/v1/contents/generations/tasks", got) | |||
| } | |||
| func TestFetchTaskUsesTianyiYunEndpoint(t *testing.T) { | |||
| adaptor := &TaskAdaptor{} | |||
| req, err := adaptor.buildFetchRequest("https://ai.ctaigw.cn", "sk-test", map[string]any{"task_id": "task-123"}) | |||
| require.NoError(t, err) | |||
| require.Equal(t, http.MethodGet, req.Method) | |||
| require.Equal(t, "https://ai.ctaigw.cn/v1/contents/generations/tasks/task-123", req.URL.String()) | |||
| require.Equal(t, "Bearer sk-test", req.Header.Get("Authorization")) | |||
| require.Equal(t, "application/json", req.Header.Get("Accept")) | |||
| } | |||
| func TestGetModelListIncludesTianyiYunSeedanceModels(t *testing.T) { | |||
| models := (&TaskAdaptor{}).GetModelList() | |||
| require.Contains(t, models, "cdance2.0-0611") | |||
| require.Contains(t, models, "cdance2.0-fast-0611") | |||
| require.Equal(t, "DoubaoVideoCompatibleTianyiYun", (&TaskAdaptor{}).GetChannelName()) | |||
| } | |||
| func TestBuildRequestBodyPreservesTianyiYunNativeFields(t *testing.T) { | |||
| adaptor := &TaskAdaptor{} | |||
| c := newTianyiYunTaskRequestContext(t, `{ | |||
| "model":"cdance2.0-0611", | |||
| "prompt":"write a clean product video", | |||
| "metadata":{ | |||
| "content":[ | |||
| {"type":"image_url","image_url":{"url":"https://example.test/input.png"},"role":"reference_image"} | |||
| ], | |||
| "ratio":"16:9", | |||
| "duration":5, | |||
| "watermark":false | |||
| } | |||
| }`) | |||
| info := &relaycommon.RelayInfo{ | |||
| ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "cdance2.0-0611"}, | |||
| } | |||
| body, err := adaptor.BuildRequestBody(c, info) | |||
| require.NoError(t, err) | |||
| data, err := io.ReadAll(body) | |||
| require.NoError(t, err) | |||
| require.Contains(t, string(data), `"model":"cdance2.0-0611"`) | |||
| require.Contains(t, string(data), `"image_url":{"url":"https://example.test/input.png"}`) | |||
| require.Contains(t, string(data), `"text":"write a clean product video"`) | |||
| require.Contains(t, string(data), `"ratio":"16:9"`) | |||
| require.Contains(t, string(data), `"duration":5`) | |||
| require.Contains(t, string(data), `"watermark":false`) | |||
| require.NotContains(t, string(data), `"seconds"`) | |||
| } | |||
| func TestBuildRequestBodyUsesMappedTianyiYunUpstreamModel(t *testing.T) { | |||
| adaptor := &TaskAdaptor{} | |||
| c := newTianyiYunTaskRequestContext(t, `{ | |||
| "model":"Doubao-Seedance-2.0", | |||
| "prompt":"write a clean product video" | |||
| }`) | |||
| info := &relaycommon.RelayInfo{ | |||
| OriginModelName: "Doubao-Seedance-2.0", | |||
| ChannelMeta: &relaycommon.ChannelMeta{ | |||
| UpstreamModelName: "cdance2.0-0611", | |||
| IsModelMapped: true, | |||
| }, | |||
| TaskRelayInfo: &relaycommon.TaskRelayInfo{}, | |||
| } | |||
| body, err := adaptor.BuildRequestBody(c, info) | |||
| require.NoError(t, err) | |||
| data, err := io.ReadAll(body) | |||
| require.NoError(t, err) | |||
| require.Contains(t, string(data), `"model":"cdance2.0-0611"`) | |||
| require.NotContains(t, string(data), `"model":"Doubao-Seedance-2.0"`) | |||
| } | |||
| func TestDoResponseRewritesIDToPublicTaskID(t *testing.T) { | |||
| adaptor := &TaskAdaptor{} | |||
| w := httptest.NewRecorder() | |||
| c, _ := gin.CreateTestContext(w) | |||
| info := &relaycommon.RelayInfo{ | |||
| OriginModelName: "cdance2.0-0611", | |||
| TaskRelayInfo: &relaycommon.TaskRelayInfo{PublicTaskID: "task_public"}, | |||
| } | |||
| resp := &http.Response{ | |||
| StatusCode: http.StatusOK, | |||
| Body: io.NopCloser(strings.NewReader(`{ | |||
| "id":"task-upstream", | |||
| "status":"queued", | |||
| "created_at":1781496040, | |||
| "model":"cdance2.0-0611" | |||
| }`)), | |||
| } | |||
| upstreamID, data, taskErr := adaptor.DoResponse(c, resp, info) | |||
| require.Nil(t, taskErr) | |||
| require.Equal(t, "task-upstream", upstreamID) | |||
| require.Contains(t, string(data), `"id":"task-upstream"`) | |||
| require.Equal(t, http.StatusOK, w.Code) | |||
| require.Contains(t, w.Body.String(), `"id":"task_public"`) | |||
| require.Contains(t, w.Body.String(), `"status":"queued"`) | |||
| require.Contains(t, w.Body.String(), `"model":"cdance2.0-0611"`) | |||
| } | |||
| func TestParseTaskResultMapsTianyiYunStatuses(t *testing.T) { | |||
| adaptor := &TaskAdaptor{} | |||
| cases := []struct { | |||
| name string | |||
| body string | |||
| wantStatus model.TaskStatus | |||
| wantURL string | |||
| wantReason string | |||
| wantUsage bool | |||
| }{ | |||
| { | |||
| name: "success", | |||
| body: `{"id":"task-ok","status":"success","content":{"video_url":"https://example.test/out.mp4"},"usage":{"completion_tokens":7,"total_tokens":11}}`, | |||
| wantStatus: model.TaskStatusSuccess, | |||
| wantURL: "https://example.test/out.mp4", | |||
| wantUsage: true, | |||
| }, | |||
| { | |||
| name: "succeeded", | |||
| body: `{"id":"task-ok","status":"succeeded","content":{"video_url":"https://example.test/out.mp4"}}`, | |||
| wantStatus: model.TaskStatusSuccess, | |||
| wantURL: "https://example.test/out.mp4", | |||
| }, | |||
| { | |||
| name: "running", | |||
| body: `{"id":"task-run","status":"running"}`, | |||
| wantStatus: model.TaskStatusInProgress, | |||
| }, | |||
| { | |||
| name: "failed", | |||
| body: `{"id":"task-fail","status":"failed","error":{"code":"BadRequest","message":"bad prompt"}}`, | |||
| wantStatus: model.TaskStatusFailure, | |||
| wantReason: "bad prompt", | |||
| }, | |||
| } | |||
| for _, tc := range cases { | |||
| t.Run(tc.name, func(t *testing.T) { | |||
| got, err := adaptor.ParseTaskResult([]byte(tc.body)) | |||
| require.NoError(t, err) | |||
| require.Equal(t, string(tc.wantStatus), got.Status) | |||
| if tc.wantURL != "" { | |||
| require.Equal(t, tc.wantURL, got.Url) | |||
| } | |||
| if tc.wantUsage { | |||
| require.Equal(t, 7, got.CompletionTokens) | |||
| require.Equal(t, 11, got.TotalTokens) | |||
| } | |||
| if tc.wantReason != "" { | |||
| require.Equal(t, tc.wantReason, got.Reason) | |||
| } | |||
| }) | |||
| } | |||
| } | |||
| func newTianyiYunTaskRequestContext(t *testing.T, body string) *gin.Context { | |||
| t.Helper() | |||
| w := httptest.NewRecorder() | |||
| c, _ := gin.CreateTestContext(w) | |||
| c.Request = httptest.NewRequest(http.MethodPost, "/api/v3/contents/generations/tasks", strings.NewReader(body)) | |||
| c.Request.Header.Set("Content-Type", "application/json") | |||
| var req relaycommon.TaskSubmitReq | |||
| require.NoError(t, common.Unmarshal([]byte(body), &req)) | |||
| relaycommon.StoreTaskRequest(c, &relaycommon.RelayInfo{}, constant.TaskActionGenerate, req) | |||
| return c | |||
| } | |||
| @@ -0,0 +1,8 @@ | |||
| package doubao_tianyiyun | |||
| var ModelList = []string{ | |||
| "cdance2.0-0611", | |||
| "cdance2.0-fast-0611", | |||
| } | |||
| var ChannelName = "DoubaoVideoCompatibleTianyiYun" | |||
| @@ -0,0 +1,391 @@ | |||
| package aiping | |||
| import ( | |||
| "bytes" | |||
| "fmt" | |||
| "io" | |||
| "net/http" | |||
| "strings" | |||
| "time" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/dto" | |||
| "github.com/QuantumNous/new-api/model" | |||
| "github.com/QuantumNous/new-api/relay/channel" | |||
| taskcommon "github.com/QuantumNous/new-api/relay/channel/task/taskcommon" | |||
| relaycommon "github.com/QuantumNous/new-api/relay/common" | |||
| "github.com/QuantumNous/new-api/service" | |||
| "github.com/gin-gonic/gin" | |||
| "github.com/pkg/errors" | |||
| "github.com/shopspring/decimal" | |||
| ) | |||
| type responsePayload struct { | |||
| Code int `json:"code"` | |||
| Message string `json:"message"` | |||
| RequestID string `json:"request_id"` | |||
| AipingID string `json:"aiping_id"` | |||
| Data struct { | |||
| TaskID string `json:"task_id"` | |||
| TaskStatus string `json:"task_status"` | |||
| TaskStatusMsg string `json:"task_status_msg"` | |||
| TaskInfo any `json:"task_info,omitempty"` | |||
| TaskResult struct { | |||
| Videos []struct { | |||
| ID string `json:"id"` | |||
| URL string `json:"url"` | |||
| WatermarkURL string `json:"watermark_url,omitempty"` | |||
| Duration string `json:"duration"` | |||
| } `json:"videos"` | |||
| } `json:"task_result"` | |||
| FinalUnitDeduction any `json:"final_unit_deduction"` | |||
| CreatedAt int64 `json:"created_at"` | |||
| UpdatedAt int64 `json:"updated_at"` | |||
| } `json:"data"` | |||
| } | |||
| type TaskAdaptor struct { | |||
| taskcommon.BaseBilling | |||
| apiKey string | |||
| baseURL string | |||
| } | |||
| func (a *TaskAdaptor) Init(info *relaycommon.RelayInfo) { | |||
| a.apiKey = info.ApiKey | |||
| a.baseURL = strings.TrimRight(info.ChannelBaseUrl, "/") | |||
| } | |||
| func (a *TaskAdaptor) ValidateRequestAndSetAction(c *gin.Context, info *relaycommon.RelayInfo) *dto.TaskError { | |||
| req, err := relaycommon.GetTaskRequest(c) | |||
| if err != nil { | |||
| return service.TaskErrorWrapperLocal(fmt.Errorf("task request not prepared"), "invalid_request", http.StatusBadRequest) | |||
| } | |||
| modelName := resolveRequestModelName(req, info) | |||
| if isOmniModel(modelName) && (info.Action == ActionText2Video || info.Action == ActionImage2Video) { | |||
| return service.TaskErrorWrapperLocal( | |||
| fmt.Errorf("model %s must use /v1/videos/omni-video", modelName), | |||
| "invalid_model_route", | |||
| http.StatusUnprocessableEntity, | |||
| ) | |||
| } | |||
| return nil | |||
| } | |||
| func (a *TaskAdaptor) BuildRequestURL(info *relaycommon.RelayInfo) (string, error) { | |||
| route, ok := FindRouteByAction(info.Action) | |||
| if !ok { | |||
| return "", fmt.Errorf("unsupported kling aiping action: %s", info.Action) | |||
| } | |||
| return a.baseURL + route.UpstreamPath, nil | |||
| } | |||
| 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 | |||
| } | |||
| func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayInfo) (io.Reader, error) { | |||
| req, err := relaycommon.GetTaskRequest(c) | |||
| if err != nil { | |||
| return nil, err | |||
| } | |||
| body := cloneMap(req.Metadata) | |||
| normalizeRequestBody(body, info) | |||
| data, err := common.Marshal(body) | |||
| if err != nil { | |||
| return nil, err | |||
| } | |||
| return bytes.NewReader(data), nil | |||
| } | |||
| func (a *TaskAdaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (*http.Response, error) { | |||
| return channel.DoTaskApiRequest(a, c, info, requestBody) | |||
| } | |||
| func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (string, []byte, *dto.TaskError) { | |||
| responseBody, err := io.ReadAll(resp.Body) | |||
| if err != nil { | |||
| return "", nil, service.TaskErrorWrapper(err, "read_response_body_failed", http.StatusInternalServerError) | |||
| } | |||
| _ = resp.Body.Close() | |||
| var parsed responsePayload | |||
| if err := common.Unmarshal(responseBody, &parsed); err != nil { | |||
| return "", nil, service.TaskErrorWrapper(errors.Wrapf(err, "body: %s", responseBody), "unmarshal_response_body_failed", http.StatusInternalServerError) | |||
| } | |||
| if parsed.Code != 0 { | |||
| msg := parsed.Message | |||
| if msg == "" { | |||
| msg = "upstream task failed" | |||
| } | |||
| return "", nil, service.TaskErrorWrapper(fmt.Errorf("%s", msg), "task_failed", http.StatusBadRequest) | |||
| } | |||
| upstreamTaskID := parsed.Data.TaskID | |||
| if strings.TrimSpace(upstreamTaskID) == "" { | |||
| return "", nil, service.TaskErrorWrapper(fmt.Errorf("task_id is empty"), "invalid_response", http.StatusInternalServerError) | |||
| } | |||
| clientPayload := map[string]any{} | |||
| if err := common.Unmarshal(responseBody, &clientPayload); err != nil { | |||
| return "", nil, service.TaskErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError) | |||
| } | |||
| sanitizeNativeTaskPayload(clientPayload, info.PublicTaskID) | |||
| c.JSON(http.StatusOK, clientPayload) | |||
| return upstreamTaskID, responseBody, nil | |||
| } | |||
| func (a *TaskAdaptor) FetchTask(baseURL, key string, body map[string]any, proxy string) (*http.Response, error) { | |||
| taskID, _ := body["task_id"].(string) | |||
| if strings.TrimSpace(taskID) == "" { | |||
| taskID, _ = body["upstream_task_id"].(string) | |||
| } | |||
| action, _ := body["action"].(string) | |||
| if strings.TrimSpace(taskID) == "" { | |||
| return nil, fmt.Errorf("invalid task_id") | |||
| } | |||
| route, ok := FindFetchRouteByAction(action) | |||
| if !ok { | |||
| return nil, fmt.Errorf("unsupported kling aiping action: %s", action) | |||
| } | |||
| url := strings.TrimRight(baseURL, "/") + strings.Replace(route.UpstreamPath, ":task_id", taskID, 1) | |||
| req, err := http.NewRequest(http.MethodGet, url, nil) | |||
| if err != nil { | |||
| return nil, err | |||
| } | |||
| req.Header.Set("Accept", "application/json") | |||
| req.Header.Set("Authorization", "Bearer "+key) | |||
| client, err := service.GetHttpClientWithProxy(proxy) | |||
| if err != nil { | |||
| return nil, fmt.Errorf("new proxy http client failed: %w", err) | |||
| } | |||
| return client.Do(req) | |||
| } | |||
| func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) { | |||
| var parsed responsePayload | |||
| if err := common.Unmarshal(respBody, &parsed); err != nil { | |||
| return nil, errors.Wrap(err, "unmarshal kling aiping task data failed") | |||
| } | |||
| info := &relaycommon.TaskInfo{ | |||
| Code: parsed.Code, | |||
| TaskID: parsed.Data.TaskID, | |||
| Reason: parsed.Data.TaskStatusMsg, | |||
| } | |||
| switch parsed.Data.TaskStatus { | |||
| case "submitted": | |||
| info.Status = model.TaskStatusSubmitted | |||
| info.Progress = taskcommon.ProgressSubmitted | |||
| case "processing": | |||
| info.Status = model.TaskStatusInProgress | |||
| info.Progress = taskcommon.ProgressInProgress | |||
| case "succeed": | |||
| info.Status = model.TaskStatusSuccess | |||
| info.Progress = taskcommon.ProgressComplete | |||
| if tokens := finalUnitDeductionTokens(parsed.Data.FinalUnitDeduction); tokens > 0 { | |||
| info.CompletionTokens = tokens | |||
| info.TotalTokens = tokens | |||
| } | |||
| case "failed": | |||
| info.Status = model.TaskStatusFailure | |||
| info.Progress = taskcommon.ProgressComplete | |||
| default: | |||
| info.Status = model.TaskStatusInProgress | |||
| info.Progress = taskcommon.ProgressInProgress | |||
| } | |||
| if len(parsed.Data.TaskResult.Videos) > 0 { | |||
| info.Url = parsed.Data.TaskResult.Videos[0].URL | |||
| } | |||
| return info, nil | |||
| } | |||
| func finalUnitDeductionTokens(value any) int { | |||
| var d decimal.Decimal | |||
| var err error | |||
| switch v := value.(type) { | |||
| case string: | |||
| if strings.TrimSpace(v) == "" { | |||
| return 0 | |||
| } | |||
| d, err = decimal.NewFromString(strings.TrimSpace(v)) | |||
| case float64: | |||
| d = decimal.NewFromFloat(v) | |||
| case int: | |||
| d = decimal.NewFromInt(int64(v)) | |||
| case int64: | |||
| d = decimal.NewFromInt(v) | |||
| default: | |||
| return 0 | |||
| } | |||
| if err != nil || !d.IsPositive() { | |||
| return 0 | |||
| } | |||
| return int(d.Mul(decimal.NewFromInt(1000000)).IntPart()) | |||
| } | |||
| func (a *TaskAdaptor) GetModelList() []string { | |||
| return Models() | |||
| } | |||
| func (a *TaskAdaptor) GetChannelName() string { | |||
| return "kling-aiping" | |||
| } | |||
| func normalizeRequestBody(body map[string]any, info *relaycommon.RelayInfo) { | |||
| delete(body, "action_control") | |||
| delete(body, "uid") | |||
| delete(body, "create_at") | |||
| delete(body, "_standard_model") | |||
| modelName := stringValue(body["model_name"]) | |||
| if modelName == "" { | |||
| modelName = stringValue(body["model"]) | |||
| } | |||
| if info != nil && info.IsModelMapped && info.UpstreamModelName != "" { | |||
| modelName = info.UpstreamModelName | |||
| } | |||
| if modelName != "" { | |||
| body["model_name"] = canonicalAipingModelName(modelName) | |||
| } | |||
| delete(body, "model") | |||
| duration := body["duration"] | |||
| if duration == nil { | |||
| duration = body["seconds"] | |||
| } | |||
| if duration != nil { | |||
| body["duration"] = fmt.Sprint(duration) | |||
| } | |||
| delete(body, "seconds") | |||
| if imageList, ok := body["reference_images"]; ok { | |||
| if _, exists := body["image_list"]; !exists { | |||
| body["image_list"] = imageList | |||
| } | |||
| delete(body, "reference_images") | |||
| } | |||
| if imageList, ok := normalizeImageList(body["image_list"]); ok { | |||
| body["image_list"] = imageList | |||
| } | |||
| if body["watermark_info"] == nil { | |||
| return | |||
| } | |||
| } | |||
| func normalizeImageList(v any) ([]map[string]any, bool) { | |||
| items, ok := v.([]any) | |||
| if !ok { | |||
| return nil, false | |||
| } | |||
| out := make([]map[string]any, 0, len(items)) | |||
| for _, item := range items { | |||
| switch val := item.(type) { | |||
| case string: | |||
| out = append(out, map[string]any{"image": val}) | |||
| case map[string]any: | |||
| image := firstString(val, "image", "image_url", "url", "base64") | |||
| if image == "" { | |||
| out = append(out, val) | |||
| continue | |||
| } | |||
| out = append(out, map[string]any{"image": image}) | |||
| default: | |||
| out = append(out, map[string]any{"image": fmt.Sprint(val)}) | |||
| } | |||
| } | |||
| return out, true | |||
| } | |||
| func sanitizeNativeTaskPayload(payload map[string]any, publicTaskID string) { | |||
| delete(payload, "aiping_id") | |||
| data, _ := payload["data"].(map[string]any) | |||
| if data == nil { | |||
| return | |||
| } | |||
| data["task_id"] = publicTaskID | |||
| ensureWatermarkURL(data) | |||
| if _, ok := data["created_at"]; !ok { | |||
| data["created_at"] = time.Now().UnixMilli() | |||
| } | |||
| if _, ok := data["updated_at"]; !ok { | |||
| data["updated_at"] = data["created_at"] | |||
| } | |||
| } | |||
| func ensureWatermarkURL(data map[string]any) { | |||
| taskResult, _ := data["task_result"].(map[string]any) | |||
| if taskResult == nil { | |||
| return | |||
| } | |||
| videos, _ := taskResult["videos"].([]any) | |||
| for _, videoAny := range videos { | |||
| video, _ := videoAny.(map[string]any) | |||
| if video == nil { | |||
| continue | |||
| } | |||
| if _, ok := video["watermark_url"]; !ok { | |||
| video["watermark_url"] = "" | |||
| } | |||
| } | |||
| } | |||
| func cloneMap(in map[string]any) map[string]any { | |||
| out := make(map[string]any, len(in)) | |||
| for k, v := range in { | |||
| out[k] = v | |||
| } | |||
| return out | |||
| } | |||
| func stringValue(v any) string { | |||
| if s, ok := v.(string); ok { | |||
| return strings.TrimSpace(s) | |||
| } | |||
| return "" | |||
| } | |||
| func canonicalAipingModelName(modelName string) string { | |||
| trimmed := strings.TrimSpace(modelName) | |||
| switch strings.ToLower(trimmed) { | |||
| case "kling-v1": | |||
| return "Kling-V1" | |||
| case "kling-v1-6": | |||
| return "Kling-V1.6" | |||
| case "kling-v2-6": | |||
| return "Kling-V2.6" | |||
| case "kling-v3": | |||
| return "Kling-V3" | |||
| case "kling-video-o1": | |||
| return "Kling-Video-O1" | |||
| case "kling-v3-omni": | |||
| return "Kling-V3-Omni" | |||
| default: | |||
| return trimmed | |||
| } | |||
| } | |||
| func firstString(m map[string]any, keys ...string) string { | |||
| for _, key := range keys { | |||
| if s := stringValue(m[key]); s != "" { | |||
| return s | |||
| } | |||
| } | |||
| return "" | |||
| } | |||
| func resolveRequestModelName(req relaycommon.TaskSubmitReq, info *relaycommon.RelayInfo) string { | |||
| if info != nil && info.ChannelMeta != nil && info.IsModelMapped && info.UpstreamModelName != "" { | |||
| return strings.TrimSpace(info.UpstreamModelName) | |||
| } | |||
| if modelName := firstString(req.Metadata, "model_name", "model"); modelName != "" { | |||
| return modelName | |||
| } | |||
| return strings.TrimSpace(req.Model) | |||
| } | |||
| func isOmniModel(modelName string) bool { | |||
| normalized := strings.ToLower(strings.TrimSpace(modelName)) | |||
| return normalized == "kling-video-o1" || normalized == "kling-v3-omni" | |||
| } | |||
| @@ -0,0 +1,286 @@ | |||
| package aiping | |||
| import ( | |||
| "io" | |||
| "net/http" | |||
| "net/http/httptest" | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/common" | |||
| relaycommon "github.com/QuantumNous/new-api/relay/common" | |||
| "github.com/QuantumNous/new-api/service" | |||
| "github.com/gin-gonic/gin" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestRoutesUseOfficialExternalAndAipingUpstreamPaths(t *testing.T) { | |||
| route, ok := FindRoute(http.MethodPost, "/v1/videos/text2video", RouteKindSubmit) | |||
| require.True(t, ok) | |||
| require.Equal(t, ActionText2Video, route.Action) | |||
| require.Equal(t, "/api/v1/multimodal/kling/videos/text2video", route.UpstreamPath) | |||
| voice, ok := FindRoute(http.MethodPost, "/v1/general/custom-voices", RouteKindSubmit) | |||
| require.True(t, ok) | |||
| require.Equal(t, "/api/v1/multimodal/kling/general/custom-voices", voice.UpstreamPath) | |||
| presets, ok := FindRoute(http.MethodGet, "/v1/general/presets-voices", RouteKindProxy) | |||
| require.True(t, ok) | |||
| require.Equal(t, "/api/v1/multimodal/kling/general/presets-voices", presets.UpstreamPath) | |||
| deleteVoices, ok := FindRoute(http.MethodPost, "/v1/general/delete-voices", RouteKindProxy) | |||
| require.True(t, ok) | |||
| require.Equal(t, ActionDeleteVoices, deleteVoices.Action) | |||
| deleteElements, ok := FindRoute(http.MethodPost, "/v1/general/delete-advanced-elements", RouteKindProxy) | |||
| require.True(t, ok) | |||
| require.Equal(t, ActionDeleteElements, deleteElements.Action) | |||
| videoExtend, ok := FindRoute(http.MethodPost, "/v1/videos/video-extend", RouteKindSubmit) | |||
| require.True(t, ok) | |||
| require.Equal(t, ActionVideoExtend, videoExtend.Action) | |||
| require.Equal(t, "/api/v1/multimodal/kling/videos/video-extend", videoExtend.UpstreamPath) | |||
| motionControl, ok := FindRoute(http.MethodPost, "/v1/videos/motion-control", RouteKindSubmit) | |||
| require.True(t, ok) | |||
| require.Equal(t, "", motionControl.BillingModel, "motion-control BillingModel should be empty so client model_name is used") | |||
| } | |||
| func TestBuildRequestURLUsesAipingKlingPath(t *testing.T) { | |||
| adaptor := &TaskAdaptor{} | |||
| info := &relaycommon.RelayInfo{ | |||
| ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: "https://aiping.cn"}, | |||
| TaskRelayInfo: &relaycommon.TaskRelayInfo{ | |||
| Action: ActionText2Video, | |||
| }, | |||
| } | |||
| adaptor.Init(info) | |||
| url, err := adaptor.BuildRequestURL(info) | |||
| require.NoError(t, err) | |||
| require.Equal(t, "https://aiping.cn/api/v1/multimodal/kling/videos/text2video", url) | |||
| } | |||
| func TestFetchTaskUsesSingleTaskPath(t *testing.T) { | |||
| service.InitHttpClient() | |||
| var gotPath string | |||
| server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | |||
| gotPath = r.URL.Path | |||
| w.Header().Set("Content-Type", "application/json") | |||
| _, _ = w.Write([]byte(`{"code":0,"data":{"task_id":"upstream-task","task_status":"submitted"}}`)) | |||
| })) | |||
| defer server.Close() | |||
| resp, err := (&TaskAdaptor{}).FetchTask(server.URL, "key", map[string]any{ | |||
| "task_id": "upstream-task", | |||
| "action": ActionText2Video, | |||
| }, "") | |||
| require.NoError(t, err) | |||
| defer resp.Body.Close() | |||
| require.Equal(t, "/api/v1/multimodal/kling/videos/text2video/upstream-task", gotPath) | |||
| } | |||
| func TestBuildRequestBodyNormalizesOfficialCompatibleFields(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| c, _ := gin.CreateTestContext(nil) | |||
| info := &relaycommon.RelayInfo{ | |||
| ChannelMeta: &relaycommon.ChannelMeta{}, | |||
| TaskRelayInfo: &relaycommon.TaskRelayInfo{ | |||
| Action: ActionMultiImage2Video, | |||
| }, | |||
| OriginModelName: "Kling-V2.6", | |||
| } | |||
| relaycommon.StoreTaskRequest(c, info, ActionMultiImage2Video, relaycommon.TaskSubmitReq{ | |||
| Model: "Kling-V2.6", | |||
| Metadata: map[string]any{ | |||
| "model": "Kling-V1.6", | |||
| "model_name": "Kling-V2.6", | |||
| "seconds": 3, | |||
| "duration": 5, | |||
| "action_control": map[string]any{"type": "continuous"}, | |||
| "uid": "user-1", | |||
| "create_at": float64(1750000000000), | |||
| "_standard_model": "Kling-V-2-6", | |||
| "image_list": []any{ | |||
| "https://example.com/1.png", | |||
| map[string]any{"image_url": "https://example.com/2.png"}, | |||
| map[string]any{"url": "https://example.com/3.png"}, | |||
| map[string]any{"base64": "abc"}, | |||
| }, | |||
| }, | |||
| }) | |||
| body, err := (&TaskAdaptor{}).BuildRequestBody(c, info) | |||
| require.NoError(t, err) | |||
| data, err := io.ReadAll(body) | |||
| require.NoError(t, err) | |||
| var got map[string]any | |||
| require.NoError(t, common.Unmarshal(data, &got)) | |||
| require.Equal(t, "Kling-V2.6", got["model_name"]) | |||
| require.NotContains(t, got, "model") | |||
| require.Equal(t, "5", got["duration"]) | |||
| require.NotContains(t, got, "seconds") | |||
| require.NotContains(t, got, "action_control") | |||
| require.NotContains(t, got, "uid") | |||
| require.NotContains(t, got, "create_at") | |||
| require.NotContains(t, got, "_standard_model") | |||
| require.Equal(t, []any{ | |||
| map[string]any{"image": "https://example.com/1.png"}, | |||
| map[string]any{"image": "https://example.com/2.png"}, | |||
| map[string]any{"image": "https://example.com/3.png"}, | |||
| map[string]any{"image": "abc"}, | |||
| }, got["image_list"]) | |||
| } | |||
| func TestBuildRequestBodyMapsCanonicalModelNameForAiping(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| c, _ := gin.CreateTestContext(nil) | |||
| info := &relaycommon.RelayInfo{ | |||
| ChannelMeta: &relaycommon.ChannelMeta{}, | |||
| TaskRelayInfo: &relaycommon.TaskRelayInfo{ | |||
| Action: ActionText2Video, | |||
| }, | |||
| OriginModelName: "kling-v2-6", | |||
| } | |||
| relaycommon.StoreTaskRequest(c, info, ActionText2Video, relaycommon.TaskSubmitReq{ | |||
| Model: "kling-v2-6", | |||
| Metadata: map[string]any{ | |||
| "model_name": "kling-v2-6", | |||
| "prompt": "prompt", | |||
| "duration": "5", | |||
| }, | |||
| }) | |||
| body, err := (&TaskAdaptor{}).BuildRequestBody(c, info) | |||
| require.NoError(t, err) | |||
| data, err := io.ReadAll(body) | |||
| require.NoError(t, err) | |||
| var got map[string]any | |||
| require.NoError(t, common.Unmarshal(data, &got)) | |||
| require.Equal(t, "Kling-V2.6", got["model_name"]) | |||
| } | |||
| func TestBuildRequestBodyUsesMappedKlingUpstreamModel(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| c, _ := gin.CreateTestContext(nil) | |||
| info := &relaycommon.RelayInfo{ | |||
| OriginModelName: "public-kling-alias", | |||
| ChannelMeta: &relaycommon.ChannelMeta{ | |||
| UpstreamModelName: "kling-v2-6", | |||
| IsModelMapped: true, | |||
| }, | |||
| TaskRelayInfo: &relaycommon.TaskRelayInfo{ | |||
| Action: ActionText2Video, | |||
| }, | |||
| } | |||
| relaycommon.StoreTaskRequest(c, info, ActionText2Video, relaycommon.TaskSubmitReq{ | |||
| Model: "public-kling-alias", | |||
| Metadata: map[string]any{ | |||
| "model_name": "public-kling-alias", | |||
| "prompt": "prompt", | |||
| "duration": "5", | |||
| }, | |||
| }) | |||
| body, err := (&TaskAdaptor{}).BuildRequestBody(c, info) | |||
| require.NoError(t, err) | |||
| data, err := io.ReadAll(body) | |||
| require.NoError(t, err) | |||
| var got map[string]any | |||
| require.NoError(t, common.Unmarshal(data, &got)) | |||
| require.Equal(t, "Kling-V2.6", got["model_name"]) | |||
| require.NotContains(t, got, "model") | |||
| } | |||
| func TestValidateRequestRejectsOmniModelOnTextOrImageRoutes(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| c, _ := gin.CreateTestContext(nil) | |||
| info := &relaycommon.RelayInfo{ | |||
| TaskRelayInfo: &relaycommon.TaskRelayInfo{ | |||
| Action: ActionText2Video, | |||
| }, | |||
| } | |||
| relaycommon.StoreTaskRequest(c, info, ActionText2Video, relaycommon.TaskSubmitReq{ | |||
| Model: "kling-v3-omni", | |||
| Metadata: map[string]any{ | |||
| "model_name": "kling-v3-omni", | |||
| }, | |||
| }) | |||
| taskErr := (&TaskAdaptor{}).ValidateRequestAndSetAction(c, info) | |||
| require.NotNil(t, taskErr) | |||
| require.Equal(t, http.StatusUnprocessableEntity, taskErr.StatusCode) | |||
| require.Contains(t, taskErr.Message, "/v1/videos/omni-video") | |||
| } | |||
| func TestParseTaskResultConvertsFinalUnitDeductionToTokens(t *testing.T) { | |||
| taskInfo, err := (&TaskAdaptor{}).ParseTaskResult([]byte(`{ | |||
| "code": 0, | |||
| "message": "success", | |||
| "request_id": "req-1", | |||
| "data": { | |||
| "task_id": "upstream-task", | |||
| "task_status": "succeed", | |||
| "task_result": { | |||
| "videos": [ | |||
| {"id": "v1", "url": "https://example.com/video.mp4", "duration": "5"} | |||
| ] | |||
| }, | |||
| "final_unit_deduction": "1.234567" | |||
| }, | |||
| "aiping_id": "internal" | |||
| }`)) | |||
| require.NoError(t, err) | |||
| require.Equal(t, 1234567, taskInfo.CompletionTokens) | |||
| require.Equal(t, 1234567, taskInfo.TotalTokens) | |||
| } | |||
| func TestFinalUnitDeductionTokens(t *testing.T) { | |||
| cases := []struct { | |||
| name string | |||
| input any | |||
| want int | |||
| }{ | |||
| {"string decimal", "1.234567", 1234567}, | |||
| {"float64", float64(1.234567), 1234567}, | |||
| {"int", int(2), 2000000}, | |||
| {"int64", int64(3), 3000000}, | |||
| {"empty string", "", 0}, | |||
| {"whitespace string", " ", 0}, | |||
| {"zero float", float64(0), 0}, | |||
| {"negative float", float64(-1.5), 0}, | |||
| {"nil", nil, 0}, | |||
| } | |||
| for _, tc := range cases { | |||
| t.Run(tc.name, func(t *testing.T) { | |||
| require.Equal(t, tc.want, finalUnitDeductionTokens(tc.input)) | |||
| }) | |||
| } | |||
| } | |||
| func TestSanitizeNativeTaskPayloadUsesPublicTaskIDAndAddsWatermarkURL(t *testing.T) { | |||
| payload := map[string]any{ | |||
| "aiping_id": "internal", | |||
| "data": map[string]any{ | |||
| "task_id": "899333358055493641", | |||
| "task_result": map[string]any{ | |||
| "videos": []any{ | |||
| map[string]any{"id": "v1", "url": "https://example.com/video.mp4"}, | |||
| }, | |||
| }, | |||
| }, | |||
| } | |||
| sanitizeNativeTaskPayload(payload, "task_public") | |||
| require.NotContains(t, payload, "aiping_id") | |||
| data := payload["data"].(map[string]any) | |||
| require.Equal(t, "task_public", data["task_id"]) | |||
| videos := data["task_result"].(map[string]any)["videos"].([]any) | |||
| require.Equal(t, "", videos[0].(map[string]any)["watermark_url"]) | |||
| } | |||
| @@ -0,0 +1,122 @@ | |||
| package aiping | |||
| import "net/http" | |||
| const ( | |||
| ActionText2Video = "kling-text2video" | |||
| ActionImage2Video = "kling-image2video" | |||
| ActionMotionControl = "kling-motion-control" | |||
| ActionOmniVideo = "kling-omni-video" | |||
| ActionMultiImage2Video = "kling-multi-image2video" | |||
| ActionVideoExtend = "kling-video-extend" | |||
| ActionAdvancedElements = "kling-advanced-elements-create" | |||
| ActionVoicesCreate = "kling-voices-create" | |||
| ActionDeleteElements = "kling-delete-elements" | |||
| ActionDeleteVoices = "kling-voices-delete" | |||
| ) | |||
| const ( | |||
| RouteKindSubmit = "submit" | |||
| RouteKindFetch = "fetch" | |||
| RouteKindList = "list" | |||
| RouteKindProxy = "proxy" | |||
| ) | |||
| const ( | |||
| ModelKlingAdvancedElements = "kling-advanced-elements" | |||
| ModelCustomVoices = "custom-voices" | |||
| ) | |||
| type Route struct { | |||
| Method string | |||
| PathPattern string | |||
| Action string | |||
| UpstreamPath string | |||
| Kind string | |||
| BillingModel string | |||
| } | |||
| var Routes = []Route{ | |||
| // 视频能力:文生视频 | |||
| {http.MethodPost, "/v1/videos/text2video", ActionText2Video, "/api/v1/multimodal/kling/videos/text2video", RouteKindSubmit, ""}, | |||
| {http.MethodGet, "/v1/videos/text2video/:task_id", ActionText2Video, "/api/v1/multimodal/kling/videos/text2video/:task_id", RouteKindFetch, ""}, | |||
| {http.MethodGet, "/v1/videos/text2video", ActionText2Video, "/api/v1/multimodal/kling/videos/text2video", RouteKindList, ""}, | |||
| // 视频能力:图生视频 | |||
| {http.MethodPost, "/v1/videos/image2video", ActionImage2Video, "/api/v1/multimodal/kling/videos/image2video", RouteKindSubmit, ""}, | |||
| {http.MethodGet, "/v1/videos/image2video/:task_id", ActionImage2Video, "/api/v1/multimodal/kling/videos/image2video/:task_id", RouteKindFetch, ""}, | |||
| {http.MethodGet, "/v1/videos/image2video", ActionImage2Video, "/api/v1/multimodal/kling/videos/image2video", RouteKindList, ""}, | |||
| // 视频能力:动作控制(model_name 可指定 Kling-V3 / Kling-V2.6) | |||
| {http.MethodPost, "/v1/videos/motion-control", ActionMotionControl, "/api/v1/multimodal/kling/videos/motion-control", RouteKindSubmit, ""}, | |||
| {http.MethodGet, "/v1/videos/motion-control/:task_id", ActionMotionControl, "/api/v1/multimodal/kling/videos/motion-control/:task_id", RouteKindFetch, ""}, | |||
| {http.MethodGet, "/v1/videos/motion-control", ActionMotionControl, "/api/v1/multimodal/kling/videos/motion-control", RouteKindList, ""}, | |||
| // 视频能力:Omni / 多镜头 | |||
| {http.MethodPost, "/v1/videos/omni-video", ActionOmniVideo, "/api/v1/multimodal/kling/videos/omni-video", RouteKindSubmit, ""}, | |||
| {http.MethodGet, "/v1/videos/omni-video/:task_id", ActionOmniVideo, "/api/v1/multimodal/kling/videos/omni-video/:task_id", RouteKindFetch, ""}, | |||
| {http.MethodGet, "/v1/videos/omni-video", ActionOmniVideo, "/api/v1/multimodal/kling/videos/omni-video", RouteKindList, ""}, | |||
| // 视频能力:多图参考生视频 | |||
| {http.MethodPost, "/v1/videos/multi-image2video", ActionMultiImage2Video, "/api/v1/multimodal/kling/videos/multi-image2video", RouteKindSubmit, ""}, | |||
| {http.MethodGet, "/v1/videos/multi-image2video/:task_id", ActionMultiImage2Video, "/api/v1/multimodal/kling/videos/multi-image2video/:task_id", RouteKindFetch, ""}, | |||
| {http.MethodGet, "/v1/videos/multi-image2video", ActionMultiImage2Video, "/api/v1/multimodal/kling/videos/multi-image2video", RouteKindList, ""}, | |||
| // 视频能力:视频延长 | |||
| {http.MethodPost, "/v1/videos/video-extend", ActionVideoExtend, "/api/v1/multimodal/kling/videos/video-extend", RouteKindSubmit, ""}, | |||
| {http.MethodGet, "/v1/videos/video-extend/:task_id", ActionVideoExtend, "/api/v1/multimodal/kling/videos/video-extend/:task_id", RouteKindFetch, ""}, | |||
| {http.MethodGet, "/v1/videos/video-extend", ActionVideoExtend, "/api/v1/multimodal/kling/videos/video-extend", RouteKindList, ""}, | |||
| // 主体能力:自定义主体 | |||
| {http.MethodPost, "/v1/general/advanced-custom-elements", ActionAdvancedElements, "/api/v1/multimodal/kling/general/advanced-custom-elements", RouteKindSubmit, ModelKlingAdvancedElements}, | |||
| {http.MethodGet, "/v1/general/advanced-custom-elements/:task_id", ActionAdvancedElements, "/api/v1/multimodal/kling/general/advanced-custom-elements/:task_id", RouteKindFetch, ModelKlingAdvancedElements}, | |||
| {http.MethodGet, "/v1/general/advanced-custom-elements", "kling-advanced-elements-list", "/api/v1/multimodal/kling/general/advanced-custom-elements", RouteKindProxy, ModelKlingAdvancedElements}, | |||
| // 主体能力:官方主体列表 | |||
| {http.MethodGet, "/v1/general/advanced-presets-elements", "kling-advanced-elements-presets", "/api/v1/multimodal/kling/general/advanced-presets-elements", RouteKindProxy, ModelKlingAdvancedElements}, | |||
| // 主体能力:删除自定义主体 | |||
| {http.MethodPost, "/v1/general/delete-advanced-elements", ActionDeleteElements, "/api/v1/multimodal/kling/general/delete-advanced-elements", RouteKindProxy, ModelKlingAdvancedElements}, | |||
| // 音色能力:自定义音色 | |||
| {http.MethodPost, "/v1/general/custom-voices", ActionVoicesCreate, "/api/v1/multimodal/kling/general/custom-voices", RouteKindSubmit, ModelCustomVoices}, | |||
| {http.MethodGet, "/v1/general/custom-voices/:task_id", ActionVoicesCreate, "/api/v1/multimodal/kling/general/custom-voices/:task_id", RouteKindFetch, ModelCustomVoices}, | |||
| {http.MethodGet, "/v1/general/custom-voices", "kling-voices-list", "/api/v1/multimodal/kling/general/custom-voices", RouteKindProxy, ModelCustomVoices}, | |||
| // 音色能力:官方音色列表 | |||
| {http.MethodGet, "/v1/general/presets-voices", "kling-voices-presets", "/api/v1/multimodal/kling/general/presets-voices", RouteKindProxy, ModelCustomVoices}, | |||
| // 音色能力:删除自定义音色 | |||
| {http.MethodPost, "/v1/general/delete-voices", ActionDeleteVoices, "/api/v1/multimodal/kling/general/delete-voices", RouteKindProxy, ModelCustomVoices}, | |||
| } | |||
| func FindRoute(method, pathPattern, kind string) (Route, bool) { | |||
| for _, route := range Routes { | |||
| if route.Method == method && route.PathPattern == pathPattern && route.Kind == kind { | |||
| return route, true | |||
| } | |||
| } | |||
| return Route{}, false | |||
| } | |||
| func FindRouteByAction(action string) (Route, bool) { | |||
| for _, route := range Routes { | |||
| if route.Action == action && route.Kind == RouteKindSubmit { | |||
| return route, true | |||
| } | |||
| } | |||
| return Route{}, false | |||
| } | |||
| func FindFetchRouteByAction(action string) (Route, bool) { | |||
| for _, route := range Routes { | |||
| if route.Action == action && route.Kind == RouteKindFetch { | |||
| return route, true | |||
| } | |||
| } | |||
| return Route{}, false | |||
| } | |||
| func Models() []string { | |||
| return []string{ | |||
| "kling-v1", | |||
| "kling-v1-6", | |||
| "kling-v2-6", | |||
| "kling-v3", | |||
| "kling-video-o1", | |||
| "kling-v3-omni", | |||
| ModelKlingAdvancedElements, | |||
| ModelCustomVoices, | |||
| } | |||
| } | |||
| @@ -155,6 +155,12 @@ type RelayInfo struct { | |||
| PriceData types.PriceData | |||
| PricingConfigSnapshot *types.PricingConfig | |||
| PricingConfigSnapshotLoaded bool | |||
| PricingDecisionFrozen *types.PricingDecision | |||
| OriginPricing *types.OriginPricingSnapshot | |||
| RequireMatrixUsageBilling bool | |||
| Request dto.Request | |||
| // RequestConversionChain records request format conversions in order, e.g. | |||
| @@ -59,6 +59,14 @@ func storeTaskRequest(c *gin.Context, info *RelayInfo, action string, requestObj | |||
| info.Action = action | |||
| c.Set("task_request", requestObj) | |||
| } | |||
| func StoreTaskRequest(c *gin.Context, info *RelayInfo, action string, requestObj TaskSubmitReq) { | |||
| if info.TaskRelayInfo == nil { | |||
| info.TaskRelayInfo = &TaskRelayInfo{} | |||
| } | |||
| storeTaskRequest(c, info, action, requestObj) | |||
| } | |||
| func GetTaskRequest(c *gin.Context) (TaskSubmitReq, error) { | |||
| v, exists := c.Get("task_request") | |||
| if !exists { | |||
| @@ -86,6 +86,10 @@ func Path2RelayMode(path string) int { | |||
| relayMode = RelayModeRerank | |||
| } else if strings.HasPrefix(path, "/v1/realtime") { | |||
| relayMode = RelayModeRealtime | |||
| } else if strings.HasPrefix(path, "/api/v3/contents/generations/tasks/") { | |||
| relayMode = RelayModeVideoFetchByID | |||
| } else if strings.HasPrefix(path, "/api/v3/contents/generations/tasks") { | |||
| relayMode = RelayModeVideoSubmit | |||
| } else if strings.HasPrefix(path, "/v1beta/models") || strings.HasPrefix(path, "/v1/models") { | |||
| relayMode = RelayModeGemini | |||
| } else if strings.HasPrefix(path, "/mj") { | |||
| @@ -14,6 +14,40 @@ import ( | |||
| "github.com/gorilla/websocket" | |||
| ) | |||
| const ( | |||
| streamBodyWrittenKey = "stream_body_written" | |||
| streamDataWrittenKey = "stream_data_written" | |||
| ) | |||
| func MarkStreamBodyWritten(c *gin.Context) { | |||
| if c == nil { | |||
| return | |||
| } | |||
| c.Set(streamBodyWrittenKey, true) | |||
| } | |||
| func StreamBodyWritten(c *gin.Context) bool { | |||
| if c == nil { | |||
| return false | |||
| } | |||
| return c.GetBool(streamBodyWrittenKey) | |||
| } | |||
| func MarkStreamDataWritten(c *gin.Context) { | |||
| if c == nil { | |||
| return | |||
| } | |||
| MarkStreamBodyWritten(c) | |||
| c.Set(streamDataWrittenKey, true) | |||
| } | |||
| func StreamDataWritten(c *gin.Context) bool { | |||
| if c == nil { | |||
| return false | |||
| } | |||
| return c.GetBool(streamDataWrittenKey) | |||
| } | |||
| func FlushWriter(c *gin.Context) (err error) { | |||
| defer func() { | |||
| if r := recover(); r != nil { | |||
| @@ -54,6 +88,19 @@ func SetEventStreamHeaders(c *gin.Context) { | |||
| c.Writer.Header().Set("X-Accel-Buffering", "no") | |||
| } | |||
| func ClearEventStreamHeadersIfNotWritten(c *gin.Context) { | |||
| if c == nil || c.Writer == nil || StreamBodyWritten(c) { | |||
| return | |||
| } | |||
| c.Set("event_stream_headers_set", false) | |||
| header := c.Writer.Header() | |||
| header.Del("Content-Type") | |||
| header.Del("Cache-Control") | |||
| header.Del("Connection") | |||
| header.Del("Transfer-Encoding") | |||
| header.Del("X-Accel-Buffering") | |||
| } | |||
| func ClaudeData(c *gin.Context, resp dto.ClaudeResponse) error { | |||
| jsonData, err := common.Marshal(resp) | |||
| if err != nil { | |||
| @@ -61,6 +108,7 @@ func ClaudeData(c *gin.Context, resp dto.ClaudeResponse) error { | |||
| } else { | |||
| c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("event: %s\n", resp.Type)}) | |||
| c.Render(-1, common.CustomEvent{Data: "data: " + string(jsonData)}) | |||
| MarkStreamDataWritten(c) | |||
| } | |||
| _ = FlushWriter(c) | |||
| return nil | |||
| @@ -69,12 +117,14 @@ func ClaudeData(c *gin.Context, resp dto.ClaudeResponse) error { | |||
| func ClaudeChunkData(c *gin.Context, resp dto.ClaudeResponse, data string) { | |||
| c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("event: %s\n", resp.Type)}) | |||
| c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("data: %s\n", data)}) | |||
| MarkStreamDataWritten(c) | |||
| _ = FlushWriter(c) | |||
| } | |||
| func ResponseChunkData(c *gin.Context, resp dto.ResponsesStreamResponse, data string) { | |||
| c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("event: %s\n", resp.Type)}) | |||
| c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("data: %s", data)}) | |||
| MarkStreamDataWritten(c) | |||
| _ = FlushWriter(c) | |||
| } | |||
| @@ -88,6 +138,7 @@ func StringData(c *gin.Context, str string) error { | |||
| } | |||
| c.Render(-1, common.CustomEvent{Data: "data: " + str}) | |||
| MarkStreamDataWritten(c) | |||
| return FlushWriter(c) | |||
| } | |||
| @@ -103,6 +154,7 @@ func PingData(c *gin.Context) error { | |||
| if _, err := c.Writer.Write([]byte(": PING\n\n")); err != nil { | |||
| return fmt.Errorf("write ping data failed: %w", err) | |||
| } | |||
| MarkStreamBodyWritten(c) | |||
| return FlushWriter(c) | |||
| } | |||
| @@ -0,0 +1,185 @@ | |||
| package helper | |||
| import ( | |||
| "fmt" | |||
| "math" | |||
| "strconv" | |||
| "strings" | |||
| "github.com/QuantumNous/new-api/common" | |||
| ) | |||
| func init() { | |||
| registerDerivedPricingFunc("megapixels", derivedMegapixels) | |||
| registerDerivedPricingFunc("aspect_ratio", derivedAspectRatio) | |||
| registerDerivedPricingFunc("input_mode", derivedInputMode) | |||
| registerDerivedPricingFunc("video_input", derivedVideoInput) | |||
| } | |||
| func derivedMegapixels(req, resp map[string]any) (any, error) { | |||
| resolution, ok, _ := getNestedValue(req, "resolution") | |||
| if !ok || !hasNonEmptyValue(resolution) { | |||
| resolution, ok, _ = getNestedValue(req, "options.size") | |||
| } | |||
| if !ok || !hasNonEmptyValue(resolution) { | |||
| return nil, nil | |||
| } | |||
| width, height, err := parseResolution(fmt.Sprintf("%v", resolution)) | |||
| if err != nil { | |||
| return nil, err | |||
| } | |||
| return float64(width*height) / 1_000_000, nil | |||
| } | |||
| func derivedAspectRatio(req, resp map[string]any) (any, error) { | |||
| resolution, ok, _ := getNestedValue(req, "resolution") | |||
| if !ok || !hasNonEmptyValue(resolution) { | |||
| resolution, ok, _ = getNestedValue(req, "options.size") | |||
| } | |||
| if !ok || !hasNonEmptyValue(resolution) { | |||
| return nil, nil | |||
| } | |||
| width, height, err := parseResolution(fmt.Sprintf("%v", resolution)) | |||
| if err != nil { | |||
| return nil, err | |||
| } | |||
| if height == 0 { | |||
| return nil, fmt.Errorf("resolution height must be greater than zero") | |||
| } | |||
| gcd := intGCD(width, height) | |||
| return fmt.Sprintf("%d:%d", width/gcd, height/gcd), nil | |||
| } | |||
| func derivedInputMode(req, resp map[string]any) (any, error) { | |||
| if hasAnyNonEmpty(req, "video", "video_url", "reference_video", "first_frame_video") { | |||
| return "v2v", nil | |||
| } | |||
| if hasAnyNonEmpty(req, "image", "image_url", "image[]", "images", "first_frame_image", "last_frame_image") { | |||
| return "i2v", nil | |||
| } | |||
| return "t2v", nil | |||
| } | |||
| func derivedVideoInput(req, resp map[string]any) (any, error) { | |||
| if hasAnyNonEmpty(req, "video", "video_url", "reference_video", "first_frame_video") { | |||
| return true, nil | |||
| } | |||
| metadata, ok, _ := getNestedValue(req, "metadata") | |||
| if !ok || !hasNonEmptyValue(metadata) { | |||
| return false, nil | |||
| } | |||
| metadataMap, err := normalizeMetadataMap(metadata) | |||
| if err != nil { | |||
| return nil, err | |||
| } | |||
| return metadataContainsVideoInput(metadataMap), nil | |||
| } | |||
| func normalizeMetadataMap(raw any) (map[string]any, error) { | |||
| switch value := raw.(type) { | |||
| case map[string]any: | |||
| return value, nil | |||
| case string: | |||
| if strings.TrimSpace(value) == "" { | |||
| return nil, nil | |||
| } | |||
| metadata := map[string]any{} | |||
| if err := common.Unmarshal([]byte(value), &metadata); err != nil { | |||
| return nil, err | |||
| } | |||
| return metadata, nil | |||
| default: | |||
| return nil, nil | |||
| } | |||
| } | |||
| func metadataContainsVideoInput(metadata map[string]any) bool { | |||
| if metadata == nil { | |||
| return false | |||
| } | |||
| content, ok := metadata["content"].([]any) | |||
| if !ok { | |||
| return false | |||
| } | |||
| for _, item := range content { | |||
| m, ok := item.(map[string]any) | |||
| if !ok { | |||
| continue | |||
| } | |||
| if typ, _ := m["type"].(string); typ == "video_url" { | |||
| return true | |||
| } | |||
| if value, ok := m["video_url"]; ok && hasNonEmptyValue(value) { | |||
| return true | |||
| } | |||
| } | |||
| return false | |||
| } | |||
| func hasNonEmptyValue(v any) bool { | |||
| switch value := v.(type) { | |||
| case nil: | |||
| return false | |||
| case string: | |||
| return strings.TrimSpace(value) != "" | |||
| case []any: | |||
| return len(value) > 0 | |||
| case []string: | |||
| return len(value) > 0 | |||
| case map[string]any: | |||
| return len(value) > 0 | |||
| default: | |||
| return true | |||
| } | |||
| } | |||
| func hasAnyNonEmpty(req map[string]any, keys ...string) bool { | |||
| for _, key := range keys { | |||
| value, ok, _ := getNestedValue(req, key) | |||
| if ok && hasNonEmptyValue(value) { | |||
| return true | |||
| } | |||
| } | |||
| return false | |||
| } | |||
| func parseResolution(raw string) (int, int, error) { | |||
| normalized := strings.ToLower(strings.TrimSpace(raw)) | |||
| if strings.HasSuffix(normalized, "p") { | |||
| height, err := strconv.Atoi(strings.TrimSuffix(normalized, "p")) | |||
| if err != nil || height <= 0 { | |||
| return 0, 0, fmt.Errorf("invalid resolution %q", raw) | |||
| } | |||
| width := int(math.Round(float64(height) * 16 / 9)) | |||
| return width, height, nil | |||
| } | |||
| parts := strings.FieldsFunc(normalized, func(r rune) bool { | |||
| return r == 'x' || r == '*' || r == '×' | |||
| }) | |||
| if len(parts) != 2 { | |||
| return 0, 0, fmt.Errorf("invalid resolution %q", raw) | |||
| } | |||
| width, err := strconv.Atoi(strings.TrimSpace(parts[0])) | |||
| if err != nil || width <= 0 { | |||
| return 0, 0, fmt.Errorf("invalid resolution width %q", raw) | |||
| } | |||
| height, err := strconv.Atoi(strings.TrimSpace(parts[1])) | |||
| if err != nil || height <= 0 { | |||
| return 0, 0, fmt.Errorf("invalid resolution height %q", raw) | |||
| } | |||
| return width, height, nil | |||
| } | |||
| func intGCD(a, b int) int { | |||
| for b != 0 { | |||
| a, b = b, a%b | |||
| } | |||
| if a < 0 { | |||
| return -a | |||
| } | |||
| if a == 0 { | |||
| return 1 | |||
| } | |||
| return a | |||
| } | |||
| @@ -0,0 +1,87 @@ | |||
| package helper | |||
| import ( | |||
| "testing" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestDerivedInputMode(t *testing.T) { | |||
| mode, err := derivedInputMode(map[string]any{}, nil) | |||
| require.NoError(t, err) | |||
| require.Equal(t, "t2v", mode) | |||
| mode, err = derivedInputMode(map[string]any{"image": "frame.png"}, nil) | |||
| require.NoError(t, err) | |||
| require.Equal(t, "i2v", mode) | |||
| mode, err = derivedInputMode(map[string]any{"image": "frame.png", "video": "clip.mp4"}, nil) | |||
| require.NoError(t, err) | |||
| require.Equal(t, "v2v", mode) | |||
| } | |||
| func TestDerivedVideoInput(t *testing.T) { | |||
| value, err := derivedVideoInput(map[string]any{}, nil) | |||
| require.NoError(t, err) | |||
| require.Equal(t, false, value) | |||
| value, err = derivedVideoInput(map[string]any{"video_url": "https://example.test/input.mp4"}, nil) | |||
| require.NoError(t, err) | |||
| require.Equal(t, true, value) | |||
| value, err = derivedVideoInput(map[string]any{ | |||
| "metadata": map[string]any{ | |||
| "content": []any{ | |||
| map[string]any{ | |||
| "type": "video_url", | |||
| "video_url": map[string]any{ | |||
| "url": "https://example.test/input.mp4", | |||
| }, | |||
| }, | |||
| }, | |||
| }, | |||
| }, nil) | |||
| require.NoError(t, err) | |||
| require.Equal(t, true, value) | |||
| value, err = derivedVideoInput(map[string]any{ | |||
| "metadata": `{"content":[{"type":"video_url","video_url":{"url":"https://example.test/input.mp4"}}]}`, | |||
| }, nil) | |||
| require.NoError(t, err) | |||
| require.Equal(t, true, value) | |||
| value, err = derivedVideoInput(map[string]any{ | |||
| "metadata": map[string]any{ | |||
| "content": []any{ | |||
| map[string]any{ | |||
| "type": "image_url", | |||
| "image_url": map[string]any{ | |||
| "url": "https://example.test/input.png", | |||
| }, | |||
| }, | |||
| }, | |||
| }, | |||
| }, nil) | |||
| require.NoError(t, err) | |||
| require.Equal(t, false, value) | |||
| } | |||
| func TestDerivedResolutionHelpers(t *testing.T) { | |||
| mp, err := derivedMegapixels(map[string]any{"resolution": "1280x720"}, nil) | |||
| require.NoError(t, err) | |||
| require.InDelta(t, 0.9216, mp, 0.0001) | |||
| ratio, err := derivedAspectRatio(map[string]any{"options": map[string]any{"size": "1280x720"}}, nil) | |||
| require.NoError(t, err) | |||
| require.Equal(t, "16:9", ratio) | |||
| } | |||
| func TestHasNonEmptyValue(t *testing.T) { | |||
| require.False(t, hasNonEmptyValue("")) | |||
| require.False(t, hasNonEmptyValue(" ")) | |||
| require.False(t, hasNonEmptyValue(nil)) | |||
| require.False(t, hasNonEmptyValue([]any{})) | |||
| require.True(t, hasNonEmptyValue([]any{"x"})) | |||
| require.True(t, hasNonEmptyValue("x")) | |||
| require.True(t, hasNonEmptyValue(0)) | |||
| } | |||
| @@ -0,0 +1,251 @@ | |||
| package helper | |||
| import ( | |||
| "fmt" | |||
| "io" | |||
| "mime/multipart" | |||
| "strconv" | |||
| "strings" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/types" | |||
| "github.com/gin-gonic/gin" | |||
| ) | |||
| type derivedPricingFunc func(req, resp map[string]any) (any, error) | |||
| var derivedPricingFuncs = map[string]derivedPricingFunc{} | |||
| type ResolvedDimensionValues struct { | |||
| Values map[string]any | |||
| Missing map[string]bool | |||
| RawInput map[string]any | |||
| } | |||
| type DimensionResolver struct { | |||
| c *gin.Context | |||
| request map[string]any | |||
| response map[string]any | |||
| } | |||
| func NewDimensionResolver(c *gin.Context) (*DimensionResolver, error) { | |||
| contentType := c.Request.Header.Get("Content-Type") | |||
| req := map[string]any{} | |||
| switch { | |||
| case strings.HasPrefix(contentType, "application/json"): | |||
| storage, err := common.GetBodyStorage(c) | |||
| if err != nil { | |||
| return nil, err | |||
| } | |||
| body, err := storage.Bytes() | |||
| if err != nil { | |||
| return nil, err | |||
| } | |||
| if len(strings.TrimSpace(string(body))) > 0 { | |||
| if err := common.Unmarshal(body, &req); err != nil { | |||
| return nil, err | |||
| } | |||
| } | |||
| if err := resetRequestBody(c, storage); err != nil { | |||
| return nil, err | |||
| } | |||
| case strings.Contains(contentType, gin.MIMEMultipartPOSTForm): | |||
| form, err := common.ParseMultipartFormReusable(c) | |||
| if err != nil { | |||
| return nil, err | |||
| } | |||
| req = requestMapFromMultipart(form) | |||
| default: | |||
| return nil, fmt.Errorf("unsupported content-type for pricing dimensions: %s", contentType) | |||
| } | |||
| return &DimensionResolver{ | |||
| c: c, | |||
| request: req, | |||
| response: map[string]any{}, | |||
| }, nil | |||
| } | |||
| func (r *DimensionResolver) ResolveDimensions(dims []types.PricingDimension) (*ResolvedDimensionValues, error) { | |||
| resolved := &ResolvedDimensionValues{ | |||
| Values: make(map[string]any, len(dims)), | |||
| Missing: make(map[string]bool), | |||
| RawInput: types.CloneMapAny(r.request), | |||
| } | |||
| for _, dim := range dims { | |||
| value, exists, err := r.resolveSource(dim.Source) | |||
| if err != nil { | |||
| return nil, err | |||
| } | |||
| if !exists || !hasNonEmptyValue(value) { | |||
| if dim.Optional { | |||
| resolved.Values[dim.Key] = dim.Default | |||
| continue | |||
| } | |||
| resolved.Values[dim.Key] = nil | |||
| resolved.Missing[dim.Key] = true | |||
| continue | |||
| } | |||
| converted, err := convertPricingDimensionValue(value, dim.Type) | |||
| if err != nil { | |||
| return nil, fmt.Errorf("dimension %q: %w", dim.Key, err) | |||
| } | |||
| resolved.Values[dim.Key] = converted | |||
| } | |||
| return resolved, nil | |||
| } | |||
| func IsAllowedPricingSource(source string) bool { | |||
| if strings.HasPrefix(source, "response.") { | |||
| return false | |||
| } | |||
| if strings.HasPrefix(source, "request.") { | |||
| return len(strings.TrimPrefix(source, "request.")) > 0 | |||
| } | |||
| if _, ok := derivedPricingFuncs[strings.TrimPrefix(source, "derived.")]; ok && strings.HasPrefix(source, "derived.") { | |||
| return true | |||
| } | |||
| return false | |||
| } | |||
| func registerDerivedPricingFunc(name string, fn derivedPricingFunc) { | |||
| derivedPricingFuncs[name] = fn | |||
| } | |||
| func (r *DimensionResolver) resolveSource(source string) (any, bool, error) { | |||
| if strings.HasPrefix(source, "response.") { | |||
| return nil, false, fmt.Errorf("pricing source %q is unsupported in schema v1", source) | |||
| } | |||
| if strings.HasPrefix(source, "request.") { | |||
| return getNestedValue(r.request, strings.TrimPrefix(source, "request.")) | |||
| } | |||
| if strings.HasPrefix(source, "derived.") { | |||
| name := strings.TrimPrefix(source, "derived.") | |||
| fn, ok := derivedPricingFuncs[name] | |||
| if !ok { | |||
| return nil, false, fmt.Errorf("derived pricing source %q is not registered", source) | |||
| } | |||
| value, err := fn(r.request, r.response) | |||
| if err != nil { | |||
| return nil, false, err | |||
| } | |||
| return value, hasNonEmptyValue(value), nil | |||
| } | |||
| return nil, false, fmt.Errorf("pricing source %q is not allowed", source) | |||
| } | |||
| func getNestedValue(src map[string]any, path string) (any, bool, error) { | |||
| if path == "" { | |||
| return nil, false, nil | |||
| } | |||
| var cur any = src | |||
| for _, part := range strings.Split(path, ".") { | |||
| m, ok := cur.(map[string]any) | |||
| if !ok { | |||
| return nil, false, nil | |||
| } | |||
| cur, ok = m[part] | |||
| if !ok { | |||
| return nil, false, nil | |||
| } | |||
| } | |||
| return cur, true, nil | |||
| } | |||
| func convertPricingDimensionValue(value any, typ string) (any, error) { | |||
| switch typ { | |||
| case "string": | |||
| if s, ok := value.(string); ok { | |||
| return s, nil | |||
| } | |||
| return fmt.Sprintf("%v", value), nil | |||
| case "number": | |||
| switch v := value.(type) { | |||
| case float64: | |||
| return v, nil | |||
| case float32: | |||
| return float64(v), nil | |||
| case int: | |||
| return float64(v), nil | |||
| case int64: | |||
| return float64(v), nil | |||
| case int32: | |||
| return float64(v), nil | |||
| case string: | |||
| f, err := strconv.ParseFloat(strings.TrimSpace(v), 64) | |||
| if err != nil { | |||
| return nil, fmt.Errorf("value %q is not a number", v) | |||
| } | |||
| return f, nil | |||
| default: | |||
| return nil, fmt.Errorf("value is not a number") | |||
| } | |||
| case "boolean": | |||
| switch v := value.(type) { | |||
| case bool: | |||
| return v, nil | |||
| case string: | |||
| b, err := strconv.ParseBool(strings.TrimSpace(v)) | |||
| if err != nil { | |||
| return nil, fmt.Errorf("value %q is not a boolean", v) | |||
| } | |||
| return b, nil | |||
| default: | |||
| return nil, fmt.Errorf("value is not a boolean") | |||
| } | |||
| default: | |||
| return nil, fmt.Errorf("unknown dimension type %q", typ) | |||
| } | |||
| } | |||
| func requestMapFromMultipart(form *multipart.Form) map[string]any { | |||
| req := map[string]any{} | |||
| for key, vals := range form.Value { | |||
| if len(vals) == 1 { | |||
| assignNestedValue(req, key, vals[0]) | |||
| continue | |||
| } | |||
| copied := append([]string(nil), vals...) | |||
| assignNestedValue(req, key, copied) | |||
| } | |||
| for key, files := range form.File { | |||
| if len(files) == 1 { | |||
| assignNestedValue(req, key, files[0].Filename) | |||
| continue | |||
| } | |||
| names := make([]string, 0, len(files)) | |||
| for _, file := range files { | |||
| names = append(names, file.Filename) | |||
| } | |||
| assignNestedValue(req, key, names) | |||
| } | |||
| return req | |||
| } | |||
| func assignNestedValue(dst map[string]any, key string, value any) { | |||
| parts := strings.Split(key, ".") | |||
| cur := dst | |||
| for i, part := range parts { | |||
| if i == len(parts)-1 { | |||
| cur[part] = value | |||
| return | |||
| } | |||
| next, ok := cur[part].(map[string]any) | |||
| if !ok { | |||
| next = map[string]any{} | |||
| cur[part] = next | |||
| } | |||
| cur = next | |||
| } | |||
| } | |||
| func resetRequestBody(c *gin.Context, storage common.BodyStorage) error { | |||
| if _, err := storage.Seek(0, io.SeekStart); err != nil { | |||
| return err | |||
| } | |||
| c.Request.Body = io.NopCloser(storage) | |||
| return nil | |||
| } | |||
| @@ -0,0 +1,147 @@ | |||
| package helper | |||
| import ( | |||
| "bytes" | |||
| "mime/multipart" | |||
| "net/http" | |||
| "net/http/httptest" | |||
| "strings" | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/types" | |||
| "github.com/gin-gonic/gin" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestDimensionResolver_JSONBody(t *testing.T) { | |||
| c := newPricingJSONContext(`{"resolution":"720P","duration":"10","options":{"size":"1280x720"},"image_url":"https://example.test/a.png"}`) | |||
| resolver, err := NewDimensionResolver(c) | |||
| require.NoError(t, err) | |||
| values, err := resolver.ResolveDimensions([]types.PricingDimension{ | |||
| {Key: "resolution", Source: "request.resolution", Type: "string"}, | |||
| {Key: "size", Source: "request.options.size", Type: "string"}, | |||
| {Key: "duration", Source: "request.duration", Type: "number"}, | |||
| {Key: "mode", Source: "derived.input_mode", Type: "string"}, | |||
| }) | |||
| require.NoError(t, err) | |||
| require.Equal(t, "720P", values.Values["resolution"]) | |||
| require.Equal(t, "1280x720", values.Values["size"]) | |||
| require.Equal(t, 10.0, values.Values["duration"]) | |||
| require.Equal(t, "i2v", values.Values["mode"]) | |||
| require.Empty(t, values.Missing) | |||
| } | |||
| func TestDimensionResolver_DerivedInputModePriority(t *testing.T) { | |||
| c := newPricingJSONContext(`{"image_url":"https://example.test/a.png","video_url":"https://example.test/a.mp4"}`) | |||
| resolver, err := NewDimensionResolver(c) | |||
| require.NoError(t, err) | |||
| values, err := resolver.ResolveDimensions([]types.PricingDimension{ | |||
| {Key: "mode", Source: "derived.input_mode", Type: "string"}, | |||
| }) | |||
| require.NoError(t, err) | |||
| require.Equal(t, "v2v", values.Values["mode"]) | |||
| } | |||
| func TestDimensionResolver_DerivedVideoInputFromMetadataContent(t *testing.T) { | |||
| c := newPricingJSONContext(`{ | |||
| "model":"seedance-2", | |||
| "prompt":"test", | |||
| "metadata":{ | |||
| "content":[ | |||
| {"type":"video_url","video_url":{"url":"https://example.test/input.mp4"}} | |||
| ] | |||
| } | |||
| }`) | |||
| resolver, err := NewDimensionResolver(c) | |||
| require.NoError(t, err) | |||
| values, err := resolver.ResolveDimensions([]types.PricingDimension{ | |||
| {Key: "video_input", Source: "derived.video_input", Type: "boolean"}, | |||
| }) | |||
| require.NoError(t, err) | |||
| require.Equal(t, true, values.Values["video_input"]) | |||
| } | |||
| func TestDimensionResolver_MissingRequiredValue(t *testing.T) { | |||
| c := newPricingJSONContext(`{"resolution":"720P"}`) | |||
| resolver, err := NewDimensionResolver(c) | |||
| require.NoError(t, err) | |||
| values, err := resolver.ResolveDimensions([]types.PricingDimension{ | |||
| {Key: "duration", Source: "request.duration", Type: "number"}, | |||
| }) | |||
| require.NoError(t, err) | |||
| require.True(t, values.Missing["duration"]) | |||
| require.Nil(t, values.Values["duration"]) | |||
| } | |||
| func TestDimensionResolver_OptionalDefault(t *testing.T) { | |||
| c := newPricingJSONContext(`{"resolution":"720P"}`) | |||
| resolver, err := NewDimensionResolver(c) | |||
| require.NoError(t, err) | |||
| values, err := resolver.ResolveDimensions([]types.PricingDimension{ | |||
| {Key: "duration", Source: "request.duration", Type: "number", Optional: true, Default: 5.0}, | |||
| }) | |||
| require.NoError(t, err) | |||
| require.Equal(t, 5.0, values.Values["duration"]) | |||
| require.Empty(t, values.Missing) | |||
| } | |||
| func TestDimensionResolver_MultipartFileInputMode(t *testing.T) { | |||
| body := &bytes.Buffer{} | |||
| writer := multipart.NewWriter(body) | |||
| require.NoError(t, writer.WriteField("resolution", "720P")) | |||
| part, err := writer.CreateFormFile("image", "frame.png") | |||
| require.NoError(t, err) | |||
| _, err = part.Write([]byte("fake")) | |||
| require.NoError(t, err) | |||
| require.NoError(t, writer.Close()) | |||
| c := newPricingContext(body.String(), writer.FormDataContentType()) | |||
| resolver, err := NewDimensionResolver(c) | |||
| require.NoError(t, err) | |||
| values, err := resolver.ResolveDimensions([]types.PricingDimension{ | |||
| {Key: "resolution", Source: "request.resolution", Type: "string"}, | |||
| {Key: "mode", Source: "derived.input_mode", Type: "string"}, | |||
| }) | |||
| require.NoError(t, err) | |||
| require.Equal(t, "720P", values.Values["resolution"]) | |||
| require.Equal(t, "i2v", values.Values["mode"]) | |||
| } | |||
| func TestDimensionResolver_UnsupportedContentType(t *testing.T) { | |||
| c := newPricingContext("resolution=720P", "application/x-www-form-urlencoded") | |||
| _, err := NewDimensionResolver(c) | |||
| require.ErrorContains(t, err, "unsupported content-type") | |||
| } | |||
| func TestDimensionResolver_ResponseSourceUnsupported(t *testing.T) { | |||
| c := newPricingJSONContext(`{"resolution":"720P"}`) | |||
| resolver, err := NewDimensionResolver(c) | |||
| require.NoError(t, err) | |||
| _, err = resolver.ResolveDimensions([]types.PricingDimension{ | |||
| {Key: "tokens", Source: "response.usage.total_tokens", Type: "number"}, | |||
| }) | |||
| require.ErrorContains(t, err, "unsupported in schema v1") | |||
| } | |||
| func newPricingJSONContext(body string) *gin.Context { | |||
| return newPricingContext(body, "application/json") | |||
| } | |||
| func newPricingContext(body string, contentType string) *gin.Context { | |||
| gin.SetMode(gin.TestMode) | |||
| w := httptest.NewRecorder() | |||
| c, _ := gin.CreateTestContext(w) | |||
| req, _ := http.NewRequest(http.MethodPost, "/v1/video", strings.NewReader(body)) | |||
| req.Header.Set("Content-Type", contentType) | |||
| c.Request = req | |||
| return c | |||
| } | |||
| @@ -0,0 +1,79 @@ | |||
| package helper | |||
| import ( | |||
| "errors" | |||
| commonpkg "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/constant" | |||
| ) | |||
| type MatrixUsageCapability struct { | |||
| BillingModelName string | |||
| UpstreamModelName string | |||
| ChannelType int | |||
| } | |||
| var matrixUsageCapabilities = []MatrixUsageCapability{ | |||
| {BillingModelName: "seedance-2", UpstreamModelName: "seedance-2", ChannelType: constant.ChannelTypeDoubaoVideo}, | |||
| {BillingModelName: "doubao-video", UpstreamModelName: "doubao-video", ChannelType: constant.ChannelTypeDoubaoVideo}, | |||
| {BillingModelName: "doubao-video-seedance", UpstreamModelName: "doubao-video-seedance", ChannelType: constant.ChannelTypeDoubaoVideo}, | |||
| {BillingModelName: "doubao-seedance-2-0-260128", UpstreamModelName: "doubao-seedance-2-0-260128", ChannelType: constant.ChannelTypeDoubaoVideo}, | |||
| {BillingModelName: "doubao-seedance-2-0-fast-260128", UpstreamModelName: "doubao-seedance-2-0-fast-260128", ChannelType: constant.ChannelTypeDoubaoVideo}, | |||
| {BillingModelName: "doubao-seedance-2-0-260128", UpstreamModelName: "doubao-seedance-2-0-260128", ChannelType: constant.ChannelTypeDoubaoVideoCompatibleAiping}, | |||
| {BillingModelName: "doubao-seedance-2-0-fast-260128", UpstreamModelName: "doubao-seedance-2-0-fast-260128", ChannelType: constant.ChannelTypeDoubaoVideoCompatibleAiping}, | |||
| {BillingModelName: "cdance2.0-0611", UpstreamModelName: "cdance2.0-0611", ChannelType: constant.ChannelTypeDoubaoVideoCompatibleTianyiYun}, | |||
| {BillingModelName: "cdance2.0-fast-0611", UpstreamModelName: "cdance2.0-fast-0611", ChannelType: constant.ChannelTypeDoubaoVideoCompatibleTianyiYun}, | |||
| // Keep in sync with SupportsAnyMatrixUsageBillingModel in setting/ratio_setting/model_pricing.go | |||
| {BillingModelName: "kling-v1", UpstreamModelName: "kling-v1", ChannelType: constant.ChannelTypeKlingAiping}, | |||
| {BillingModelName: "kling-v1", UpstreamModelName: "Kling-V1", ChannelType: constant.ChannelTypeKlingAiping}, | |||
| {BillingModelName: "kling-v1-6", UpstreamModelName: "kling-v1-6", ChannelType: constant.ChannelTypeKlingAiping}, | |||
| {BillingModelName: "kling-v1-6", UpstreamModelName: "Kling-V1.6", ChannelType: constant.ChannelTypeKlingAiping}, | |||
| {BillingModelName: "kling-v2-6", UpstreamModelName: "kling-v2-6", ChannelType: constant.ChannelTypeKlingAiping}, | |||
| {BillingModelName: "kling-v2-6", UpstreamModelName: "Kling-V2.6", ChannelType: constant.ChannelTypeKlingAiping}, | |||
| {BillingModelName: "kling-v3", UpstreamModelName: "kling-v3", ChannelType: constant.ChannelTypeKlingAiping}, | |||
| {BillingModelName: "kling-v3", UpstreamModelName: "Kling-V3", ChannelType: constant.ChannelTypeKlingAiping}, | |||
| {BillingModelName: "kling-video-o1", UpstreamModelName: "kling-video-o1", ChannelType: constant.ChannelTypeKlingAiping}, | |||
| {BillingModelName: "kling-video-o1", UpstreamModelName: "Kling-Video-O1", ChannelType: constant.ChannelTypeKlingAiping}, | |||
| {BillingModelName: "kling-v3-omni", UpstreamModelName: "kling-v3-omni", ChannelType: constant.ChannelTypeKlingAiping}, | |||
| {BillingModelName: "kling-v3-omni", UpstreamModelName: "Kling-V3-Omni", ChannelType: constant.ChannelTypeKlingAiping}, | |||
| } | |||
| func SupportsMatrixUsageBilling(modelName string, channelType int, mappedModel string) bool { | |||
| if mappedModel == "" { | |||
| mappedModel = modelName | |||
| } | |||
| for _, capability := range matrixUsageCapabilities { | |||
| if capability.BillingModelName == modelName && | |||
| capability.UpstreamModelName == mappedModel && | |||
| capability.ChannelType == channelType { | |||
| return true | |||
| } | |||
| } | |||
| return false | |||
| } | |||
| func ResolvePricingModelMapping(modelName string, modelMapping string) (string, error) { | |||
| if modelMapping == "" || modelMapping == "{}" { | |||
| return modelName, nil | |||
| } | |||
| modelMap := map[string]string{} | |||
| if err := commonpkg.Unmarshal([]byte(modelMapping), &modelMap); err != nil { | |||
| return "", err | |||
| } | |||
| current := modelName | |||
| visited := map[string]bool{current: true} | |||
| for { | |||
| next := modelMap[current] | |||
| if next == "" { | |||
| return current, nil | |||
| } | |||
| if visited[next] { | |||
| if next == current { | |||
| return current, nil | |||
| } | |||
| return "", errors.New("model_mapping_contains_cycle") | |||
| } | |||
| visited[next] = true | |||
| current = next | |||
| } | |||
| } | |||
| @@ -0,0 +1,62 @@ | |||
| package helper | |||
| import ( | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/constant" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestSupportsMatrixUsageBilling(t *testing.T) { | |||
| require.True(t, SupportsMatrixUsageBilling("seedance-2", constant.ChannelTypeDoubaoVideo, "seedance-2")) | |||
| require.False(t, SupportsMatrixUsageBilling("seedance-2", constant.ChannelTypeOpenAI, "seedance-2")) | |||
| require.False(t, SupportsMatrixUsageBilling("unknown", constant.ChannelTypeDoubaoVideo, "unknown")) | |||
| } | |||
| func TestSupportsMatrixUsageBilling_DoubaoSeedanceOfficialModelNames(t *testing.T) { | |||
| require.True(t, SupportsMatrixUsageBilling("doubao-seedance-2-0-260128", constant.ChannelTypeDoubaoVideo, "doubao-seedance-2-0-260128")) | |||
| require.True(t, SupportsMatrixUsageBilling("doubao-seedance-2-0-fast-260128", constant.ChannelTypeDoubaoVideo, "doubao-seedance-2-0-fast-260128")) | |||
| } | |||
| func TestSupportsMatrixUsageBilling_DoubaoAipingChannel(t *testing.T) { | |||
| require.True(t, SupportsMatrixUsageBilling("doubao-seedance-2-0-260128", constant.ChannelTypeDoubaoVideoCompatibleAiping, "doubao-seedance-2-0-260128")) | |||
| require.True(t, SupportsMatrixUsageBilling("doubao-seedance-2-0-fast-260128", constant.ChannelTypeDoubaoVideoCompatibleAiping, "doubao-seedance-2-0-fast-260128")) | |||
| } | |||
| func TestSupportsMatrixUsageBilling_TianyiYunSeedanceChannel(t *testing.T) { | |||
| require.True(t, SupportsMatrixUsageBilling("cdance2.0-0611", constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "cdance2.0-0611")) | |||
| require.True(t, SupportsMatrixUsageBilling("cdance2.0-fast-0611", constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "cdance2.0-fast-0611")) | |||
| require.False(t, SupportsMatrixUsageBilling("cdance2.0-0611", constant.ChannelTypeDoubaoVideoCompatibleAiping, "cdance2.0-0611")) | |||
| } | |||
| func TestSupportsMatrixUsageBilling_KlingAipingVideoModels(t *testing.T) { | |||
| cases := []struct { | |||
| billingModel string | |||
| upstreamModel string | |||
| }{ | |||
| {"kling-v1", "Kling-V1"}, | |||
| {"kling-v1-6", "Kling-V1.6"}, | |||
| {"kling-v2-6", "Kling-V2.6"}, | |||
| {"kling-v3", "Kling-V3"}, | |||
| {"kling-video-o1", "Kling-Video-O1"}, | |||
| {"kling-v3-omni", "Kling-V3-Omni"}, | |||
| } | |||
| for _, tc := range cases { | |||
| require.True(t, SupportsMatrixUsageBilling(tc.billingModel, constant.ChannelTypeKlingAiping, tc.billingModel), tc.billingModel) | |||
| require.True(t, SupportsMatrixUsageBilling(tc.billingModel, constant.ChannelTypeKlingAiping, tc.upstreamModel), tc.billingModel) | |||
| } | |||
| } | |||
| func TestResolvePricingModelMapping(t *testing.T) { | |||
| mapped, err := ResolvePricingModelMapping("seedance-2", `{"seedance-2":"doubao-video-seedance"}`) | |||
| require.NoError(t, err) | |||
| require.Equal(t, "doubao-video-seedance", mapped) | |||
| mapped, err = ResolvePricingModelMapping("seedance-2", `{"seedance-2":"a","a":"b"}`) | |||
| require.NoError(t, err) | |||
| require.Equal(t, "b", mapped) | |||
| _, err = ResolvePricingModelMapping("seedance-2", `{"seedance-2":"a","a":"seedance-2"}`) | |||
| require.ErrorContains(t, err, "model_mapping_contains_cycle") | |||
| } | |||
| @@ -2,8 +2,10 @@ package helper | |||
| import ( | |||
| "fmt" | |||
| "net/http" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/dto" | |||
| "github.com/QuantumNous/new-api/logger" | |||
| relaycommon "github.com/QuantumNous/new-api/relay/common" | |||
| "github.com/QuantumNous/new-api/setting/operation_setting" | |||
| @@ -124,9 +126,57 @@ func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens | |||
| } | |||
| // ModelPriceHelperPerCall 按次计费的 PriceHelper (MJ、Task) | |||
| func ModelPriceHelperPerCall(c *gin.Context, info *relaycommon.RelayInfo) types.PriceData { | |||
| func ModelPriceHelperPerCall(c *gin.Context, info *relaycommon.RelayInfo) (types.PriceData, *dto.TaskError) { | |||
| groupRatioInfo := HandleGroupRatio(c, info) | |||
| if info.PricingDecisionFrozen != nil { | |||
| priceData := priceDataFromDecision(info.PricingDecisionFrozen) | |||
| info.PriceData = priceData | |||
| return priceData, nil | |||
| } | |||
| if info.OriginPricing != nil { | |||
| var ( | |||
| decision *types.PricingDecision | |||
| taskErr *dto.TaskError | |||
| ) | |||
| if info.OriginPricing.BillingMode == types.BillingModeMatrix { | |||
| decision, taskErr = buildRemixMatrixDecision(info.OriginPricing, groupRatioInfo) | |||
| } else { | |||
| decision, taskErr = buildRemixLegacyDecision(info.OriginPricing, groupRatioInfo) | |||
| } | |||
| if taskErr != nil { | |||
| return types.PriceData{}, taskErr | |||
| } | |||
| info.PricingDecisionFrozen = decision | |||
| priceData := priceDataFromDecision(decision) | |||
| info.PriceData = priceData | |||
| return priceData, nil | |||
| } | |||
| if info.PricingConfigSnapshotLoaded && info.PricingConfigSnapshot != nil { | |||
| resolver, err := NewDimensionResolver(c) | |||
| if err != nil { | |||
| return types.PriceData{}, pricingDimensionTaskError(err) | |||
| } | |||
| resolved, err := resolver.ResolveDimensions(info.PricingConfigSnapshot.Dimensions) | |||
| if err != nil { | |||
| return types.PriceData{}, pricingDimensionTaskError(err) | |||
| } | |||
| result, taskErr := lookupPricingConfig(info.PricingConfigSnapshot, resolved) | |||
| if taskErr != nil { | |||
| return types.PriceData{}, taskErr | |||
| } | |||
| decision := buildMatchedDecision(info.PricingConfigSnapshot, result, groupRatioInfo) | |||
| if result.Mode != types.PricingModeMatched { | |||
| decision = buildFallbackDecision(info.PricingConfigSnapshot, result, groupRatioInfo) | |||
| } | |||
| info.PricingDecisionFrozen = decision | |||
| priceData := priceDataFromDecision(decision) | |||
| info.PriceData = priceData | |||
| return priceData, nil | |||
| } | |||
| modelPrice, success := ratio_setting.GetModelPrice(info.OriginModelName, true) | |||
| if !success { | |||
| defaultPrice, ok := ratio_setting.GetDefaultModelPriceMap()[info.OriginModelName] | |||
| @@ -147,12 +197,26 @@ func ModelPriceHelperPerCall(c *gin.Context, info *relaycommon.RelayInfo) types. | |||
| } | |||
| priceData := types.PriceData{ | |||
| FreeModel: freeModel, | |||
| ModelPrice: modelPrice, | |||
| Quota: quota, | |||
| GroupRatioInfo: groupRatioInfo, | |||
| FreeModel: freeModel, | |||
| ModelPrice: modelPrice, | |||
| Quota: quota, | |||
| GroupRatioInfo: groupRatioInfo, | |||
| } | |||
| decision := buildLegacyDecision(modelPrice, groupRatioInfo, freeModel, quota) | |||
| priceData.PricingDecision = decision | |||
| info.PricingDecisionFrozen = decision | |||
| info.PriceData = priceData | |||
| return priceData, nil | |||
| } | |||
| func pricingDimensionTaskError(err error) *dto.TaskError { | |||
| return &dto.TaskError{ | |||
| Code: "pricing_dimension_resolve_failed", | |||
| Message: err.Error(), | |||
| StatusCode: http.StatusBadRequest, | |||
| LocalError: true, | |||
| Error: err, | |||
| } | |||
| return priceData | |||
| } | |||
| func ContainPriceOrRatio(modelName string) bool { | |||
| @@ -116,7 +116,6 @@ func TestGlobalRatiosFilledWhenUseRatioMode(t *testing.T) { | |||
| assert.Equal(t, 2.0, priceData.ImageRatio) | |||
| } | |||
| func TestHandleGroupRatio_DefaultGroup(t *testing.T) { | |||
| setupUsePriceSwitchTest(t) | |||
| c := buildTestContext(t) | |||
| @@ -305,3 +304,62 @@ func TestModelPriceHelper_FreeGroup(t *testing.T) { | |||
| // cleanup | |||
| require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1}`)) | |||
| } | |||
| func TestModelPriceHelperPerCall_MatrixMatchFreezesDecision(t *testing.T) { | |||
| setupUsePriceSwitchTest(t) | |||
| c := newPricingJSONContext(`{"resolution":"720P","duration":10}`) | |||
| info := &relaycommon.RelayInfo{ | |||
| OriginModelName: "hailuo-video", | |||
| UsingGroup: "default", | |||
| UserGroup: "default", | |||
| PricingConfigSnapshotLoaded: true, | |||
| PricingConfigSnapshot: testPricingLookupConfig(types.PricingFallbackReject), | |||
| } | |||
| priceData, taskErr := ModelPriceHelperPerCall(c, info) | |||
| require.Nil(t, taskErr) | |||
| require.Equal(t, int(0.08*common.QuotaPerUnit), priceData.Quota) | |||
| require.NotNil(t, priceData.PricingDecision) | |||
| require.Same(t, info.PricingDecisionFrozen, priceData.PricingDecision) | |||
| require.Equal(t, types.BillingModeMatrix, priceData.PricingDecision.BillingMode) | |||
| info.PricingConfigSnapshot.Table[0]["price"] = 1.5 | |||
| next, taskErr := ModelPriceHelperPerCall(c, info) | |||
| require.Nil(t, taskErr) | |||
| require.Equal(t, priceData.Quota, next.Quota) | |||
| } | |||
| func TestModelPriceHelperPerCall_MatrixReject(t *testing.T) { | |||
| setupUsePriceSwitchTest(t) | |||
| c := newPricingJSONContext(`{"resolution":"1080P","duration":15}`) | |||
| info := &relaycommon.RelayInfo{ | |||
| OriginModelName: "hailuo-video", | |||
| UsingGroup: "default", | |||
| UserGroup: "default", | |||
| PricingConfigSnapshotLoaded: true, | |||
| PricingConfigSnapshot: testPricingLookupConfig(types.PricingFallbackReject), | |||
| } | |||
| priceData, taskErr := ModelPriceHelperPerCall(c, info) | |||
| require.NotNil(t, taskErr) | |||
| require.Equal(t, "pricing_no_match", taskErr.Code) | |||
| require.Nil(t, info.PricingDecisionFrozen) | |||
| require.Zero(t, priceData.Quota) | |||
| } | |||
| func TestModelPriceHelperPerCall_LegacyDecision(t *testing.T) { | |||
| setupUsePriceSwitchTest(t) | |||
| c := buildTestContext(t) | |||
| info := &relaycommon.RelayInfo{ | |||
| OriginModelName: testModelUsePriceSwitch, | |||
| UsingGroup: "default", | |||
| UserGroup: "default", | |||
| } | |||
| priceData, taskErr := ModelPriceHelperPerCall(c, info) | |||
| require.Nil(t, taskErr) | |||
| require.Equal(t, int(0.5*common.QuotaPerUnit), priceData.Quota) | |||
| require.NotNil(t, priceData.PricingDecision) | |||
| require.Equal(t, types.PricingModeLegacy, priceData.PricingDecision.Mode) | |||
| require.Same(t, info.PricingDecisionFrozen, priceData.PricingDecision) | |||
| } | |||
| @@ -0,0 +1,292 @@ | |||
| package helper | |||
| import ( | |||
| "fmt" | |||
| "net/http" | |||
| "strings" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/dto" | |||
| "github.com/QuantumNous/new-api/types" | |||
| ) | |||
| type pricingLookupResult struct { | |||
| Mode types.PricingMode | |||
| PriceUSD float64 | |||
| Snapshot map[string]any | |||
| } | |||
| func lookupPricingConfig(cfg *types.PricingConfig, resolved *ResolvedDimensionValues) (*pricingLookupResult, *dto.TaskError) { | |||
| if cfg == nil { | |||
| return nil, localPricingTaskError("pricing config is required", "pricing_config_required") | |||
| } | |||
| var best types.PricingRow | |||
| bestWildcardCount := len(cfg.Dimensions) + 1 | |||
| bestCount := 0 | |||
| for _, row := range cfg.Table { | |||
| wildcards, ok := rowMatchesDimensions(cfg.Dimensions, row, resolved.Values) | |||
| if !ok { | |||
| continue | |||
| } | |||
| if wildcards < bestWildcardCount { | |||
| best = row | |||
| bestWildcardCount = wildcards | |||
| bestCount = 1 | |||
| continue | |||
| } | |||
| if wildcards == bestWildcardCount { | |||
| bestCount++ | |||
| } | |||
| } | |||
| if best != nil && bestCount == 1 { | |||
| price, _ := pricingNumberValue(best["price"]) | |||
| return &pricingLookupResult{ | |||
| Mode: types.PricingModeMatched, | |||
| PriceUSD: price, | |||
| Snapshot: buildPricingSnapshot(cfg, resolved, price, types.PricingModeMatched, ""), | |||
| }, nil | |||
| } | |||
| if bestCount > 1 { | |||
| return fallbackPricingConfig(cfg, resolved) | |||
| } | |||
| return fallbackPricingConfig(cfg, resolved) | |||
| } | |||
| func buildMatchedDecision(cfg *types.PricingConfig, result *pricingLookupResult, group types.GroupRatioInfo) *types.PricingDecision { | |||
| return buildMatrixDecision(cfg, result, group) | |||
| } | |||
| func buildFallbackDecision(cfg *types.PricingConfig, result *pricingLookupResult, group types.GroupRatioInfo) *types.PricingDecision { | |||
| return buildMatrixDecision(cfg, result, group) | |||
| } | |||
| func buildLegacyDecision(modelPrice float64, group types.GroupRatioInfo, free bool, quota int) *types.PricingDecision { | |||
| return &types.PricingDecision{ | |||
| Mode: types.PricingModeLegacy, | |||
| PriceUSD: modelPrice, | |||
| Quota: quota, | |||
| GroupRatioInfo: group, | |||
| FreeModel: free, | |||
| LegacyModelPrice: modelPrice, | |||
| UsePrice: true, | |||
| PerCallBilling: true, | |||
| Finalized: true, | |||
| } | |||
| } | |||
| func buildRemixMatrixDecision(origin *types.OriginPricingSnapshot, group types.GroupRatioInfo) (*types.PricingDecision, *dto.TaskError) { | |||
| if origin == nil { | |||
| return nil, localPricingTaskError("origin pricing snapshot is required", "invalid_pricing_snapshot") | |||
| } | |||
| billingUnit := origin.BillingUnit | |||
| if billingUnit == "" { | |||
| billingUnit = types.BillingUnitPerCall | |||
| } | |||
| result := &pricingLookupResult{ | |||
| Mode: types.PricingModeRemix, | |||
| PriceUSD: origin.PriceUSD, | |||
| Snapshot: types.CloneMapAny(origin.Snapshot), | |||
| } | |||
| cfg := &types.PricingConfig{BillingUnit: billingUnit} | |||
| decision := buildMatrixDecision(cfg, result, group) | |||
| decision.Mode = types.PricingModeRemix | |||
| decision.TokenUnitPriceUSD = origin.TokenUnitPriceUSD | |||
| decision.OtherRatios = types.CloneRatios(origin.OtherRatios) | |||
| decision.PerCallBilling = origin.PerCallBilling | |||
| if billingUnit == types.BillingUnitPer1MTokens { | |||
| decision.PriceUSD = origin.TokenUnitPriceUSD | |||
| decision.Quota = minimumUsagePreconsumeQuota(origin.TokenUnitPriceUSD, group.GroupRatio) | |||
| } | |||
| return decision, nil | |||
| } | |||
| func buildRemixLegacyDecision(origin *types.OriginPricingSnapshot, group types.GroupRatioInfo) (*types.PricingDecision, *dto.TaskError) { | |||
| if origin == nil { | |||
| return nil, localPricingTaskError("origin legacy pricing snapshot is required", "invalid_legacy_pricing_snapshot") | |||
| } | |||
| quota := int(origin.PriceUSD * common.QuotaPerUnit * group.GroupRatio) | |||
| return &types.PricingDecision{ | |||
| Mode: types.PricingModeRemix, | |||
| PriceUSD: origin.PriceUSD, | |||
| Quota: quota, | |||
| GroupRatioInfo: group, | |||
| LegacyModelPrice: origin.PriceUSD, | |||
| OtherRatios: types.CloneRatios(origin.OtherRatios), | |||
| UsePrice: true, | |||
| PerCallBilling: origin.PerCallBilling, | |||
| Finalized: true, | |||
| }, nil | |||
| } | |||
| func buildMatrixDecision(cfg *types.PricingConfig, result *pricingLookupResult, group types.GroupRatioInfo) *types.PricingDecision { | |||
| billingUnit := cfg.BillingUnit | |||
| if billingUnit == "" { | |||
| billingUnit = types.BillingUnitPerCall | |||
| } | |||
| decision := &types.PricingDecision{ | |||
| Mode: result.Mode, | |||
| BillingMode: types.BillingModeMatrix, | |||
| BillingUnit: billingUnit, | |||
| GroupRatioInfo: group, | |||
| Snapshot: types.CloneMapAny(result.Snapshot), | |||
| UsePrice: true, | |||
| Finalized: true, | |||
| } | |||
| if billingUnit == types.BillingUnitPer1MTokens { | |||
| decision.TokenUnitPriceUSD = result.PriceUSD | |||
| decision.PriceUSD = result.PriceUSD | |||
| decision.Quota = minimumUsagePreconsumeQuota(result.PriceUSD, group.GroupRatio) | |||
| decision.PerCallBilling = false | |||
| return decision | |||
| } | |||
| decision.PriceUSD = result.PriceUSD | |||
| decision.Quota = int(result.PriceUSD * common.QuotaPerUnit * group.GroupRatio) | |||
| decision.PerCallBilling = true | |||
| if group.GroupRatio == 0 || result.PriceUSD == 0 { | |||
| decision.FreeModel = true | |||
| decision.Quota = 0 | |||
| } | |||
| return decision | |||
| } | |||
| func fallbackPricingConfig(cfg *types.PricingConfig, resolved *ResolvedDimensionValues) (*pricingLookupResult, *dto.TaskError) { | |||
| switch cfg.Fallback.Strategy { | |||
| case types.PricingFallbackReject: | |||
| return nil, localPricingTaskError("pricing dimensions did not match any row", "pricing_no_match") | |||
| case types.PricingFallbackMax: | |||
| price := maxPricingTablePrice(cfg.Table) | |||
| return &pricingLookupResult{ | |||
| Mode: types.PricingModeFallbackMax, | |||
| PriceUSD: price, | |||
| Snapshot: buildPricingSnapshot(cfg, resolved, price, types.PricingModeFallbackMax, cfg.Fallback.Strategy), | |||
| }, nil | |||
| case types.PricingFallbackDefault: | |||
| price := cfg.Fallback.DefaultPrice | |||
| return &pricingLookupResult{ | |||
| Mode: types.PricingModeFallbackDefault, | |||
| PriceUSD: price, | |||
| Snapshot: buildPricingSnapshot(cfg, resolved, price, types.PricingModeFallbackDefault, cfg.Fallback.Strategy), | |||
| }, nil | |||
| default: | |||
| return nil, localPricingTaskError(fmt.Sprintf("unsupported fallback strategy %q", cfg.Fallback.Strategy), "invalid_pricing_config") | |||
| } | |||
| } | |||
| func buildPricingSnapshot(cfg *types.PricingConfig, resolved *ResolvedDimensionValues, price float64, mode types.PricingMode, fallbackStrategy string) map[string]any { | |||
| snapshot := make(map[string]any, len(cfg.Dimensions)+4) | |||
| for _, dim := range cfg.Dimensions { | |||
| snapshot[dim.Key] = resolved.Values[dim.Key] | |||
| } | |||
| snapshot["price"] = price | |||
| snapshot["billing_unit"] = cfg.BillingUnit | |||
| snapshot["pricing_mode"] = string(mode) | |||
| if fallbackStrategy != "" { | |||
| snapshot["fallback_strategy"] = fallbackStrategy | |||
| } | |||
| return snapshot | |||
| } | |||
| func rowMatchesDimensions(dimensions []types.PricingDimension, row types.PricingRow, values map[string]any) (int, bool) { | |||
| wildcards := 0 | |||
| for _, dim := range dimensions { | |||
| rv := row[dim.Key] | |||
| if s, ok := rv.(string); ok && s == "*" { | |||
| wildcards++ | |||
| continue | |||
| } | |||
| if !samePricingLiteral(rv, values[dim.Key], dim.Type) { | |||
| return 0, false | |||
| } | |||
| } | |||
| return wildcards, true | |||
| } | |||
| func maxPricingTablePrice(rows []types.PricingRow) float64 { | |||
| var max float64 | |||
| for i, row := range rows { | |||
| price, ok := pricingNumberValue(row["price"]) | |||
| if !ok { | |||
| continue | |||
| } | |||
| if i == 0 || price > max { | |||
| max = price | |||
| } | |||
| } | |||
| return max | |||
| } | |||
| func minimumUsagePreconsumeQuota(tokenUnitPriceUSD float64, groupRatio float64) int { | |||
| if tokenUnitPriceUSD == 0 || groupRatio == 0 { | |||
| return 0 | |||
| } | |||
| if common.PreConsumedQuota <= 0 { | |||
| return 0 | |||
| } | |||
| quota := int(float64(common.PreConsumedQuota) * tokenUnitPriceUSD / 1_000_000 * common.QuotaPerUnit * groupRatio) | |||
| if quota <= 0 { | |||
| return 1 | |||
| } | |||
| return quota | |||
| } | |||
| func priceDataFromDecision(decision *types.PricingDecision) types.PriceData { | |||
| if decision == nil { | |||
| return types.PriceData{} | |||
| } | |||
| return types.PriceData{ | |||
| FreeModel: decision.FreeModel, | |||
| ModelPrice: decision.PriceUSD, | |||
| Quota: decision.Quota, | |||
| GroupRatioInfo: decision.GroupRatioInfo, | |||
| OtherRatios: types.CloneRatios(decision.OtherRatios), | |||
| UsePrice: decision.UsePrice, | |||
| PricingDecision: decision, | |||
| } | |||
| } | |||
| func pricingNumberValue(v any) (float64, bool) { | |||
| switch n := v.(type) { | |||
| case float64: | |||
| return n, true | |||
| case float32: | |||
| return float64(n), true | |||
| case int: | |||
| return float64(n), true | |||
| case int64: | |||
| return float64(n), true | |||
| case int32: | |||
| return float64(n), true | |||
| default: | |||
| return 0, false | |||
| } | |||
| } | |||
| func samePricingLiteral(a, b any, dimensionType string) bool { | |||
| af, aok := pricingNumberValue(a) | |||
| bf, bok := pricingNumberValue(b) | |||
| if aok && bok { | |||
| return af == bf | |||
| } | |||
| if dimensionType == "string" { | |||
| return strings.EqualFold(strings.TrimSpace(fmt.Sprintf("%v", a)), strings.TrimSpace(fmt.Sprintf("%v", b))) | |||
| } | |||
| return fmt.Sprintf("%v", a) == fmt.Sprintf("%v", b) | |||
| } | |||
| func localPricingTaskError(message string, code string) *dto.TaskError { | |||
| err := fmt.Errorf("%s", message) | |||
| return &dto.TaskError{ | |||
| Code: code, | |||
| Message: message, | |||
| StatusCode: http.StatusBadRequest, | |||
| LocalError: true, | |||
| Error: err, | |||
| } | |||
| } | |||
| @@ -0,0 +1,147 @@ | |||
| package helper | |||
| import ( | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/types" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestPricingLookup_MatchSpecificity(t *testing.T) { | |||
| cfg := testPricingLookupConfig(types.PricingFallbackReject) | |||
| resolved := &ResolvedDimensionValues{Values: map[string]any{"resolution": "720P", "duration": 10.0}} | |||
| result, taskErr := lookupPricingConfig(cfg, resolved) | |||
| require.Nil(t, taskErr) | |||
| require.Equal(t, types.PricingModeMatched, result.Mode) | |||
| require.Equal(t, 0.08, result.PriceUSD) | |||
| require.Equal(t, "720P", result.Snapshot["resolution"]) | |||
| require.Equal(t, 10.0, result.Snapshot["duration"]) | |||
| } | |||
| func TestPricingLookup_StringDimensionMatchesCaseInsensitive(t *testing.T) { | |||
| cfg := testPricingLookupConfig(types.PricingFallbackReject) | |||
| resolved := &ResolvedDimensionValues{Values: map[string]any{"resolution": "720p", "duration": 10.0}} | |||
| result, taskErr := lookupPricingConfig(cfg, resolved) | |||
| require.Nil(t, taskErr) | |||
| require.Equal(t, types.PricingModeMatched, result.Mode) | |||
| require.Equal(t, 0.08, result.PriceUSD) | |||
| require.Equal(t, "720p", result.Snapshot["resolution"]) | |||
| } | |||
| func TestPricingLookup_WildcardFallbackMatch(t *testing.T) { | |||
| cfg := testPricingLookupConfig(types.PricingFallbackReject) | |||
| resolved := &ResolvedDimensionValues{Values: map[string]any{"resolution": "480P", "duration": 10.0}} | |||
| result, taskErr := lookupPricingConfig(cfg, resolved) | |||
| require.Nil(t, taskErr) | |||
| require.Equal(t, 0.05, result.PriceUSD) | |||
| } | |||
| func TestPricingLookup_RejectFallback(t *testing.T) { | |||
| cfg := testPricingLookupConfig(types.PricingFallbackReject) | |||
| resolved := &ResolvedDimensionValues{Values: map[string]any{"resolution": "1080P", "duration": 15.0}} | |||
| result, taskErr := lookupPricingConfig(cfg, resolved) | |||
| require.Nil(t, result) | |||
| require.NotNil(t, taskErr) | |||
| require.Equal(t, "pricing_no_match", taskErr.Code) | |||
| } | |||
| func TestPricingLookup_MaxFallback(t *testing.T) { | |||
| cfg := testPricingLookupConfig(types.PricingFallbackMax) | |||
| resolved := &ResolvedDimensionValues{Values: map[string]any{"resolution": "1080P", "duration": 15.0}} | |||
| result, taskErr := lookupPricingConfig(cfg, resolved) | |||
| require.Nil(t, taskErr) | |||
| require.Equal(t, types.PricingModeFallbackMax, result.Mode) | |||
| require.Equal(t, 0.1, result.PriceUSD) | |||
| require.Equal(t, types.PricingFallbackMax, result.Snapshot["fallback_strategy"]) | |||
| } | |||
| func TestPricingLookup_DefaultFallback(t *testing.T) { | |||
| cfg := testPricingLookupConfig(types.PricingFallbackDefault) | |||
| cfg.Fallback.DefaultPrice = 0.2 | |||
| resolved := &ResolvedDimensionValues{Values: map[string]any{"resolution": "1080P", "duration": 15.0}} | |||
| result, taskErr := lookupPricingConfig(cfg, resolved) | |||
| require.Nil(t, taskErr) | |||
| require.Equal(t, types.PricingModeFallbackDefault, result.Mode) | |||
| require.Equal(t, 0.2, result.PriceUSD) | |||
| } | |||
| func TestPricingLookup_MissingDimensionSnapshotMarshalsNull(t *testing.T) { | |||
| cfg := testPricingLookupConfig(types.PricingFallbackDefault) | |||
| cfg.Fallback.DefaultPrice = 0.2 | |||
| resolved := &ResolvedDimensionValues{ | |||
| Values: map[string]any{"resolution": "1080P", "duration": nil}, | |||
| Missing: map[string]bool{"duration": true}, | |||
| } | |||
| result, taskErr := lookupPricingConfig(cfg, resolved) | |||
| require.Nil(t, taskErr) | |||
| body, err := common.Marshal(result.Snapshot) | |||
| require.NoError(t, err) | |||
| require.Contains(t, string(body), `"duration":null`) | |||
| } | |||
| func TestPricingLookup_DecisionQuota(t *testing.T) { | |||
| cfg := testPricingLookupConfig(types.PricingFallbackReject) | |||
| result := &pricingLookupResult{Mode: types.PricingModeMatched, PriceUSD: 0.1, Snapshot: map[string]any{"price": 0.1}} | |||
| decision := buildMatchedDecision(cfg, result, types.GroupRatioInfo{GroupRatio: 0.5}) | |||
| require.Equal(t, types.BillingModeMatrix, decision.BillingMode) | |||
| require.Equal(t, types.BillingUnitPerCall, decision.BillingUnit) | |||
| require.Equal(t, int(0.1*common.QuotaPerUnit*0.5), decision.Quota) | |||
| require.True(t, decision.PerCallBilling) | |||
| } | |||
| func TestPricingLookup_UsageDecisionMinimumPreconsume(t *testing.T) { | |||
| cfg := testPricingLookupConfig(types.PricingFallbackReject) | |||
| cfg.BillingUnit = types.BillingUnitPer1MTokens | |||
| cfg.PreconsumeStrategy = types.PreconsumeStrategyMinimum | |||
| result := &pricingLookupResult{Mode: types.PricingModeMatched, PriceUSD: 2.0, Snapshot: map[string]any{"price": 2.0}} | |||
| decision := buildMatchedDecision(cfg, result, types.GroupRatioInfo{GroupRatio: 1}) | |||
| require.Equal(t, 2.0, decision.PriceUSD) | |||
| require.Equal(t, 2.0, decision.TokenUnitPriceUSD) | |||
| require.Equal(t, minimumUsagePreconsumeQuota(2.0, 1), decision.Quota) | |||
| require.False(t, decision.PerCallBilling) | |||
| } | |||
| func TestPricingLookup_RemixUsageDecisionKeepsDisplayPrice(t *testing.T) { | |||
| decision, taskErr := buildRemixMatrixDecision(&types.OriginPricingSnapshot{ | |||
| PriceUSD: 0, | |||
| TokenUnitPriceUSD: 46, | |||
| BillingUnit: types.BillingUnitPer1MTokens, | |||
| Snapshot: map[string]any{"price": 46}, | |||
| PerCallBilling: false, | |||
| }, types.GroupRatioInfo{GroupRatio: 1}) | |||
| require.Nil(t, taskErr) | |||
| require.Equal(t, 46.0, decision.PriceUSD) | |||
| require.Equal(t, 46.0, decision.TokenUnitPriceUSD) | |||
| require.Equal(t, minimumUsagePreconsumeQuota(46, 1), decision.Quota) | |||
| require.False(t, decision.PerCallBilling) | |||
| } | |||
| func testPricingLookupConfig(fallback string) *types.PricingConfig { | |||
| return &types.PricingConfig{ | |||
| SchemaVersion: types.PricingSchemaVersion, | |||
| Scope: types.PricingScopeModel, | |||
| BillingUnit: types.BillingUnitPerCall, | |||
| PreconsumeStrategy: types.PreconsumeStrategyExact, | |||
| Dimensions: []types.PricingDimension{ | |||
| {Key: "resolution", Source: "request.resolution", Type: "string"}, | |||
| {Key: "duration", Source: "request.duration", Type: "number"}, | |||
| }, | |||
| Table: []types.PricingRow{ | |||
| {"resolution": "720P", "duration": 10.0, "price": 0.08, "source": types.PricingRowSourceManual}, | |||
| {"resolution": "720P", "duration": "*", "price": 0.1, "source": types.PricingRowSourceManual}, | |||
| {"resolution": "*", "duration": 10.0, "price": 0.05, "source": types.PricingRowSourceManual}, | |||
| }, | |||
| Fallback: types.PricingFallback{Strategy: fallback}, | |||
| } | |||
| } | |||
| @@ -24,6 +24,7 @@ import ( | |||
| const ( | |||
| InitialScannerBufferSize = 64 << 10 // 64KB (64*1024) | |||
| DefaultMaxScannerBufferSize = 64 << 20 // 64MB (64*1024*1024) default SSE buffer size | |||
| DefaultStreamingTimeout = 300 * time.Second | |||
| DefaultPingInterval = 10 * time.Second | |||
| ) | |||
| @@ -48,6 +49,9 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon | |||
| }() | |||
| streamingTimeout := time.Duration(constant.StreamingTimeout) * time.Second | |||
| if streamingTimeout <= 0 { | |||
| streamingTimeout = DefaultStreamingTimeout | |||
| } | |||
| var ( | |||
| stopChan = make(chan bool, 3) // 增加缓冲区避免阻塞 | |||
| @@ -97,6 +97,24 @@ func TestStreamScannerHandler_EmptyBody(t *testing.T) { | |||
| assert.False(t, called.Load(), "handler should not be called for empty body") | |||
| } | |||
| func TestStreamScannerHandler_NonPositiveStreamingTimeoutDoesNotPanic(t *testing.T) { | |||
| oldTimeout := constant.StreamingTimeout | |||
| constant.StreamingTimeout = 0 | |||
| t.Cleanup(func() { | |||
| constant.StreamingTimeout = oldTimeout | |||
| }) | |||
| recorder := httptest.NewRecorder() | |||
| c, _ := gin.CreateTestContext(recorder) | |||
| c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) | |||
| resp := &http.Response{Body: io.NopCloser(strings.NewReader(""))} | |||
| info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{}} | |||
| require.NotPanics(t, func() { | |||
| StreamScannerHandler(c, resp, info, func(data string) bool { return true }) | |||
| }) | |||
| } | |||
| func TestStreamScannerHandler_1000Chunks(t *testing.T) { | |||
| t.Parallel() | |||
| @@ -193,7 +193,13 @@ func RelaySwapFace(c *gin.Context, info *relaycommon.RelayInfo) *dto.MidjourneyR | |||
| } | |||
| modelName := service.CovertMjpActionToModelName(constant.MjActionSwapFace) | |||
| priceData := helper.ModelPriceHelperPerCall(c, info) | |||
| priceData, taskErr := helper.ModelPriceHelperPerCall(c, info) | |||
| if taskErr != nil { | |||
| return &dto.MidjourneyResponse{ | |||
| Code: 4, | |||
| Description: taskErr.Message, | |||
| } | |||
| } | |||
| userQuota, err := model.GetUserQuota(info.UserId, false) | |||
| if err != nil { | |||
| @@ -494,7 +500,13 @@ func RelayMidjourneySubmit(c *gin.Context, relayInfo *relaycommon.RelayInfo) *dt | |||
| modelName := service.CovertMjpActionToModelName(midjRequest.Action) | |||
| priceData := helper.ModelPriceHelperPerCall(c, relayInfo) | |||
| priceData, taskErr := helper.ModelPriceHelperPerCall(c, relayInfo) | |||
| if taskErr != nil { | |||
| return &dto.MidjourneyResponse{ | |||
| Code: 4, | |||
| Description: taskErr.Message, | |||
| } | |||
| } | |||
| userQuota, err := model.GetUserQuota(relayInfo.UserId, false) | |||
| if err != nil { | |||
| @@ -32,10 +32,13 @@ import ( | |||
| "github.com/QuantumNous/new-api/relay/channel/submodel" | |||
| taskali "github.com/QuantumNous/new-api/relay/channel/task/ali" | |||
| taskdoubao "github.com/QuantumNous/new-api/relay/channel/task/doubao" | |||
| taskdoubaoaiping "github.com/QuantumNous/new-api/relay/channel/task/doubao_aiping" | |||
| taskdoubaotianyiyun "github.com/QuantumNous/new-api/relay/channel/task/doubao_tianyiyun" | |||
| taskGemini "github.com/QuantumNous/new-api/relay/channel/task/gemini" | |||
| "github.com/QuantumNous/new-api/relay/channel/task/hailuo" | |||
| taskjimeng "github.com/QuantumNous/new-api/relay/channel/task/jimeng" | |||
| "github.com/QuantumNous/new-api/relay/channel/task/kling" | |||
| klingaiping "github.com/QuantumNous/new-api/relay/channel/task/kling/aiping" | |||
| tasksora "github.com/QuantumNous/new-api/relay/channel/task/sora" | |||
| "github.com/QuantumNous/new-api/relay/channel/task/suno" | |||
| taskvertex "github.com/QuantumNous/new-api/relay/channel/task/vertex" | |||
| @@ -153,6 +156,12 @@ func GetTaskAdaptor(platform constant.TaskPlatform) channel.TaskAdaptor { | |||
| return &taskVidu.TaskAdaptor{} | |||
| case constant.ChannelTypeDoubaoVideo, constant.ChannelTypeVolcEngine: | |||
| return &taskdoubao.TaskAdaptor{} | |||
| case constant.ChannelTypeDoubaoVideoCompatibleAiping: | |||
| return &taskdoubaoaiping.TaskAdaptor{} | |||
| case constant.ChannelTypeDoubaoVideoCompatibleTianyiYun: | |||
| return &taskdoubaotianyiyun.TaskAdaptor{} | |||
| case constant.ChannelTypeKlingAiping: | |||
| return &klingaiping.TaskAdaptor{} | |||
| case constant.ChannelTypeSora, constant.ChannelTypeOpenAI: | |||
| return &tasksora.TaskAdaptor{} | |||
| case constant.ChannelTypeGemini: | |||
| @@ -0,0 +1,16 @@ | |||
| package relay | |||
| import ( | |||
| "strconv" | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/constant" | |||
| tianyiyun "github.com/QuantumNous/new-api/relay/channel/task/doubao_tianyiyun" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestGetTaskAdaptorReturnsTianyiYunSeedanceAdaptor(t *testing.T) { | |||
| adaptor := GetTaskAdaptor(constant.TaskPlatform(strconv.Itoa(constant.ChannelTypeDoubaoVideoCompatibleTianyiYun))) | |||
| require.IsType(t, &tianyiyun.TaskAdaptor{}, adaptor) | |||
| } | |||
| @@ -19,24 +19,20 @@ import ( | |||
| relayconstant "github.com/QuantumNous/new-api/relay/constant" | |||
| "github.com/QuantumNous/new-api/relay/helper" | |||
| "github.com/QuantumNous/new-api/service" | |||
| "github.com/QuantumNous/new-api/types" | |||
| "github.com/gin-gonic/gin" | |||
| ) | |||
| type TaskSubmitResult struct { | |||
| UpstreamTaskID string | |||
| TaskData []byte | |||
| Platform constant.TaskPlatform | |||
| Quota int | |||
| UpstreamTaskID string | |||
| TaskData []byte | |||
| Platform constant.TaskPlatform | |||
| Quota int | |||
| UpstreamReqJSON []byte | |||
| //PerCallPrice types.PriceData | |||
| } | |||
| // ResolveOriginTask 处理基于已有任务的提交(remix / continuation): | |||
| // 查找原始任务、从中提取模型名称、将渠道锁定到原始任务的渠道 | |||
| // (通过 info.LockedChannel,重试时复用同一渠道并轮换 key), | |||
| // 以及提取 OtherRatios(时长、分辨率)。 | |||
| // 该函数在控制器的重试循环之前调用一次,其结果通过 info 字段和上下文持久化。 | |||
| func ResolveOriginTask(c *gin.Context, info *relaycommon.RelayInfo) *dto.TaskError { | |||
| // 检测 remix action | |||
| path := c.Request.URL.Path | |||
| if strings.Contains(path, "/v1/videos/") && strings.HasSuffix(path, "/remix") { | |||
| info.Action = constant.TaskActionRemix | |||
| @@ -53,7 +49,6 @@ func ResolveOriginTask(c *gin.Context, info *relaycommon.RelayInfo) *dto.TaskErr | |||
| return nil | |||
| } | |||
| // 查找原始任务 | |||
| originTask, exist, err := model.GetByTaskId(info.UserId, info.OriginTaskID) | |||
| if err != nil { | |||
| return service.TaskErrorWrapper(err, "get_origin_task_failed", http.StatusInternalServerError) | |||
| @@ -62,7 +57,7 @@ func ResolveOriginTask(c *gin.Context, info *relaycommon.RelayInfo) *dto.TaskErr | |||
| return service.TaskErrorWrapperLocal(errors.New("task_origin_not_exist"), "task_not_exist", http.StatusBadRequest) | |||
| } | |||
| // 从原始任务推导模型名称 | |||
| // 娴犲骸甯慨瀣╂崲閸斺剝甯圭€靛吋膩閸ㄥ鎮曢敓? | |||
| if info.OriginModelName == "" { | |||
| if originTask.Properties.OriginModelName != "" { | |||
| info.OriginModelName = originTask.Properties.OriginModelName | |||
| @@ -77,7 +72,6 @@ func ResolveOriginTask(c *gin.Context, info *relaycommon.RelayInfo) *dto.TaskErr | |||
| } | |||
| } | |||
| // 锁定到原始任务的渠道(重试时复用同一渠道,轮换 key) | |||
| ch, err := model.GetChannelById(originTask.ChannelId, true) | |||
| if err != nil { | |||
| return service.TaskErrorWrapperLocal(err, "channel_not_found", http.StatusBadRequest) | |||
| @@ -87,31 +81,19 @@ func ResolveOriginTask(c *gin.Context, info *relaycommon.RelayInfo) *dto.TaskErr | |||
| } | |||
| info.LockedChannel = ch | |||
| if originTask.ChannelId != info.ChannelId { | |||
| key, _, newAPIError := ch.GetNextEnabledKey() | |||
| if newAPIError != nil { | |||
| return service.TaskErrorWrapper(newAPIError, "channel_no_available_key", newAPIError.StatusCode) | |||
| } | |||
| common.SetContextKey(c, constant.ContextKeyChannelKey, key) | |||
| common.SetContextKey(c, constant.ContextKeyChannelType, ch.Type) | |||
| common.SetContextKey(c, constant.ContextKeyChannelBaseUrl, ch.GetBaseURL()) | |||
| common.SetContextKey(c, constant.ContextKeyChannelId, originTask.ChannelId) | |||
| info.ChannelBaseUrl = ch.GetBaseURL() | |||
| info.ChannelId = originTask.ChannelId | |||
| info.ChannelType = ch.Type | |||
| info.ApiKey = key | |||
| } | |||
| // 提取 remix 参数(时长、分辨率 → OtherRatios) | |||
| if info.Action == constant.TaskActionRemix { | |||
| if originTask.PrivateData.BillingContext != nil { | |||
| // 新的 remix 逻辑:直接从原始任务的 BillingContext 中提取 OtherRatios(如果存在) | |||
| for s, f := range originTask.PrivateData.BillingContext.OtherRatios { | |||
| info.PriceData.AddOtherRatio(s, f) | |||
| bc := originTask.PrivateData.BillingContext | |||
| info.OriginPricing = &types.OriginPricingSnapshot{ | |||
| BillingMode: bc.BillingMode, | |||
| BillingUnit: bc.BillingUnit, | |||
| PriceUSD: bc.ModelPrice, | |||
| TokenUnitPriceUSD: bc.TokenUnitPriceUSD, | |||
| Snapshot: types.CloneMapAny(bc.PricingSnapshot), | |||
| OtherRatios: types.CloneRatios(bc.OtherRatios), | |||
| PerCallBilling: bc.PerCallBilling, | |||
| } | |||
| } else { | |||
| // 旧的 remix 逻辑:直接从 task data 解析 seconds 和 size(如果存在) | |||
| var taskData map[string]interface{} | |||
| _ = common.Unmarshal(originTask.Data, &taskData) | |||
| secondsStr, _ := taskData["seconds"].(string) | |||
| @@ -120,13 +102,16 @@ func ResolveOriginTask(c *gin.Context, info *relaycommon.RelayInfo) *dto.TaskErr | |||
| seconds = 4 | |||
| } | |||
| sizeStr, _ := taskData["size"].(string) | |||
| if info.PriceData.OtherRatios == nil { | |||
| info.PriceData.OtherRatios = map[string]float64{} | |||
| otherRatios := map[string]float64{ | |||
| "seconds": float64(seconds), | |||
| "size": 1, | |||
| } | |||
| info.PriceData.OtherRatios["seconds"] = float64(seconds) | |||
| info.PriceData.OtherRatios["size"] = 1 | |||
| if sizeStr == "1792x1024" || sizeStr == "1024x1792" { | |||
| info.PriceData.OtherRatios["size"] = 1.666667 | |||
| otherRatios["size"] = 1.666667 | |||
| } | |||
| info.OriginPricing = &types.OriginPricingSnapshot{ | |||
| OtherRatios: otherRatios, | |||
| PerCallBilling: false, | |||
| } | |||
| } | |||
| } | |||
| @@ -134,15 +119,9 @@ func ResolveOriginTask(c *gin.Context, info *relaycommon.RelayInfo) *dto.TaskErr | |||
| return nil | |||
| } | |||
| // RelayTaskSubmit 完成 task 提交的全部流程(每次尝试调用一次): | |||
| // 刷新渠道元数据 → 确定 platform/adaptor → 验证请求 → | |||
| // 估算计费(EstimateBilling) → 计算价格 → 预扣费(仅首次)→ | |||
| // 构建/发送/解析上游请求 → 提交后计费调整(AdjustBillingOnSubmit)。 | |||
| // 控制器负责 defer Refund 和成功后 Settle。 | |||
| func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitResult, *dto.TaskError) { | |||
| info.InitChannelMeta(c) | |||
| // 1. 确定 platform → 创建适配器 → 验证请求 | |||
| platform := constant.TaskPlatform(c.GetString("platform")) | |||
| if platform == "" { | |||
| platform = GetTaskPlatform(c) | |||
| @@ -156,39 +135,39 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe | |||
| return nil, taskErr | |||
| } | |||
| // 2. 确定模型名称 | |||
| modelName := info.OriginModelName | |||
| if modelName == "" { | |||
| modelName = service.CoverTaskActionToModelName(platform, info.Action) | |||
| } | |||
| // 2.5 应用渠道的模型映射(与同步任务对齐) | |||
| info.OriginModelName = modelName | |||
| info.UpstreamModelName = modelName | |||
| if err := helper.ModelMappedHelper(c, info, nil); err != nil { | |||
| return nil, service.TaskErrorWrapperLocal(err, "model_mapping_failed", http.StatusBadRequest) | |||
| } | |||
| // 3. 预生成公开 task ID(仅首次) | |||
| if info.PublicTaskID == "" { | |||
| info.PublicTaskID = model.GenerateTaskID() | |||
| } | |||
| // 4. 价格计算:基础模型价格 | |||
| // 4. 娴犻攱鐗哥拋锛勭暬閿涙艾鐔€绾偓濡€崇€锋禒閿嬬壐 | |||
| info.OriginModelName = modelName | |||
| info.PriceData = helper.ModelPriceHelperPerCall(c, info) | |||
| // 5. 计费估算:让适配器根据用户请求提供 OtherRatios(时长、分辨率等) | |||
| // 必须在 ModelPriceHelperPerCall 之后调用(它会重建 PriceData)。 | |||
| // ResolveOriginTask 可能已在 remix 路径中预设了 OtherRatios,此处合并。 | |||
| if estimatedRatios := adaptor.EstimateBilling(c, info); len(estimatedRatios) > 0 { | |||
| for k, v := range estimatedRatios { | |||
| info.PriceData.AddOtherRatio(k, v) | |||
| priceData, pricingErr := helper.ModelPriceHelperPerCall(c, info) | |||
| if pricingErr != nil { | |||
| return nil, pricingErr | |||
| } | |||
| info.PriceData = priceData | |||
| matrixOrOriginPricing := info.OriginPricing != nil || | |||
| (info.PricingDecisionFrozen != nil && info.PricingDecisionFrozen.BillingMode == types.BillingModeMatrix) | |||
| if !matrixOrOriginPricing { | |||
| if estimatedRatios := adaptor.EstimateBilling(c, info); len(estimatedRatios) > 0 { | |||
| for k, v := range estimatedRatios { | |||
| info.PriceData.AddOtherRatio(k, v) | |||
| } | |||
| } | |||
| } | |||
| // 6. 将 OtherRatios 应用到基础额度 | |||
| if !common.StringsContains(constant.TaskPricePatches, modelName) { | |||
| if !matrixOrOriginPricing && !common.StringsContains(constant.TaskPricePatches, modelName) { | |||
| for _, ra := range info.PriceData.OtherRatios { | |||
| if ra != 1.0 { | |||
| info.PriceData.Quota = int(float64(info.PriceData.Quota) * ra) | |||
| @@ -196,7 +175,6 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe | |||
| } | |||
| } | |||
| // 7. 预扣费(仅首次 — 重试时 info.Billing 已存在,跳过) | |||
| if info.Billing == nil && !info.PriceData.FreeModel { | |||
| info.ForcePreConsume = true | |||
| if apiErr := service.PreConsumeBilling(c, info.PriceData.Quota, info); apiErr != nil { | |||
| @@ -204,13 +182,19 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe | |||
| } | |||
| } | |||
| // 8. 构建请求体 | |||
| requestBody, err := adaptor.BuildRequestBody(c, info) | |||
| if err != nil { | |||
| return nil, service.TaskErrorWrapper(err, "build_request_failed", http.StatusInternalServerError) | |||
| } | |||
| // 9. 发送请求 | |||
| var upstreamReqBytes []byte | |||
| if requestBody != nil { | |||
| upstreamReqBytes, err = io.ReadAll(requestBody) | |||
| if err != nil { | |||
| return nil, service.TaskErrorWrapper(err, "read_request_body_failed", http.StatusInternalServerError) | |||
| } | |||
| requestBody = bytes.NewReader(upstreamReqBytes) | |||
| } | |||
| resp, err := adaptor.DoRequest(c, info, requestBody) | |||
| if err != nil { | |||
| return nil, service.TaskErrorWrapper(err, "do_request_failed", http.StatusInternalServerError) | |||
| @@ -220,7 +204,6 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe | |||
| return nil, service.TaskErrorWrapper(fmt.Errorf("%s", string(responseBody)), "fail_to_fetch_task", resp.StatusCode) | |||
| } | |||
| // 10. 返回 OtherRatios 给下游(header 必须在 DoResponse 写 body 之前设置) | |||
| otherRatios := info.PriceData.OtherRatios | |||
| if otherRatios == nil { | |||
| otherRatios = map[string]float64{} | |||
| @@ -228,41 +211,38 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe | |||
| ratiosJSON, _ := common.Marshal(otherRatios) | |||
| c.Header("X-New-Api-Other-Ratios", string(ratiosJSON)) | |||
| // 11. 解析响应 | |||
| // 11. 鐟欙絾鐎介崫宥呯安 | |||
| upstreamTaskID, taskData, taskErr := adaptor.DoResponse(c, resp, info) | |||
| if taskErr != nil { | |||
| return nil, taskErr | |||
| } | |||
| // 11. 提交后计费调整:让适配器根据上游实际返回调整 OtherRatios | |||
| finalQuota := info.PriceData.Quota | |||
| if adjustedRatios := adaptor.AdjustBillingOnSubmit(info, taskData); len(adjustedRatios) > 0 { | |||
| // 基于调整后的 ratios 重新计算 quota | |||
| finalQuota = recalcQuotaFromRatios(info, adjustedRatios) | |||
| info.PriceData.OtherRatios = adjustedRatios | |||
| info.PriceData.Quota = finalQuota | |||
| if !matrixOrOriginPricing { | |||
| if adjustedRatios := adaptor.AdjustBillingOnSubmit(info, taskData); len(adjustedRatios) > 0 { | |||
| finalQuota = recalcQuotaFromRatios(info, adjustedRatios) | |||
| info.PriceData.OtherRatios = adjustedRatios | |||
| info.PriceData.Quota = finalQuota | |||
| } | |||
| } | |||
| return &TaskSubmitResult{ | |||
| UpstreamTaskID: upstreamTaskID, | |||
| TaskData: taskData, | |||
| Platform: platform, | |||
| Quota: finalQuota, | |||
| UpstreamTaskID: upstreamTaskID, | |||
| TaskData: taskData, | |||
| Platform: platform, | |||
| Quota: finalQuota, | |||
| UpstreamReqJSON: upstreamReqBytes, | |||
| }, nil | |||
| } | |||
| // recalcQuotaFromRatios 根据 adjustedRatios 重新计算 quota。 | |||
| // 公式: baseQuota × ∏(ratio) — 其中 baseQuota 是不含 OtherRatios 的基础额度。 | |||
| func recalcQuotaFromRatios(info *relaycommon.RelayInfo, ratios map[string]float64) int { | |||
| // 从 PriceData 获取不含 OtherRatios 的基础价格 | |||
| baseQuota := info.PriceData.Quota | |||
| // 先除掉原有的 OtherRatios 恢复基础额度 | |||
| for _, ra := range info.PriceData.OtherRatios { | |||
| if ra != 1.0 && ra > 0 { | |||
| baseQuota = int(float64(baseQuota) / ra) | |||
| } | |||
| } | |||
| // 应用新的 ratios | |||
| // 鎼存梻鏁ら弬鎵畱 ratios | |||
| result := float64(baseQuota) | |||
| for _, ra := range ratios { | |||
| if ra != 1.0 { | |||
| @@ -372,13 +352,12 @@ func videoFetchByIDRespBodyBuilder(c *gin.Context) (respBody []byte, taskResp *d | |||
| isOpenAIVideoAPI := strings.HasPrefix(c.Request.RequestURI, "/v1/videos/") | |||
| // Gemini/Vertex 支持实时查询:用户 fetch 时直接从上游拉取最新状态 | |||
| // Gemini/Vertex 閺€顖涘瘮鐎圭偞妞傞弻銉嚄閿涙氨鏁ら敓?fetch 閺冨墎娲块幒銉ょ矤娑撳﹥鐖堕幏澶婂絿閺堚偓閺傛壆濮搁敓? | |||
| if realtimeResp := tryRealtimeFetch(originTask, isOpenAIVideoAPI); len(realtimeResp) > 0 { | |||
| respBody = realtimeResp | |||
| return | |||
| } | |||
| // OpenAI Video API 格式: 走各 adaptor 的 ConvertToOpenAIVideo | |||
| if isOpenAIVideoAPI { | |||
| adaptor := GetTaskAdaptor(originTask.Platform) | |||
| if adaptor == nil { | |||
| @@ -398,7 +377,6 @@ func videoFetchByIDRespBodyBuilder(c *gin.Context) (respBody []byte, taskResp *d | |||
| return | |||
| } | |||
| // 通用 TaskDto 格式 | |||
| respBody, err = common.Marshal(dto.TaskResponse[any]{ | |||
| Code: "success", | |||
| Data: TaskModel2Dto(originTask), | |||
| @@ -409,9 +387,6 @@ func videoFetchByIDRespBodyBuilder(c *gin.Context) (respBody []byte, taskResp *d | |||
| return | |||
| } | |||
| // tryRealtimeFetch 尝试从上游实时拉取 Gemini/Vertex 任务状态。 | |||
| // 仅当渠道类型为 Gemini 或 Vertex 时触发;其他渠道或出错时返回 nil。 | |||
| // 当非 OpenAI Video API 时,还会构建自定义格式的响应体。 | |||
| func tryRealtimeFetch(task *model.Task, isOpenAIVideoAPI bool) []byte { | |||
| channelModel, err := model.GetChannelById(task.ChannelId, true) | |||
| if err != nil { | |||
| @@ -451,7 +426,6 @@ func tryRealtimeFetch(task *model.Task, isOpenAIVideoAPI bool) []byte { | |||
| snap := task.Snapshot() | |||
| // 将上游最新状态更新到 task | |||
| if ti.Status != "" { | |||
| task.Status = model.TaskStatus(ti.Status) | |||
| } | |||
| @@ -459,11 +433,9 @@ func tryRealtimeFetch(task *model.Task, isOpenAIVideoAPI bool) []byte { | |||
| task.Progress = ti.Progress | |||
| } | |||
| if strings.HasPrefix(ti.Url, "data:") { | |||
| // data: URI — kept in Data, not ResultURL | |||
| } else if ti.Url != "" { | |||
| task.PrivateData.ResultURL = ti.Url | |||
| } else if task.Status == model.TaskStatusSuccess { | |||
| // No URL from adaptor — construct proxy URL using public task ID | |||
| task.PrivateData.ResultURL = taskcommon.BuildProxyURL(task.TaskID) | |||
| } | |||
| @@ -471,12 +443,10 @@ func tryRealtimeFetch(task *model.Task, isOpenAIVideoAPI bool) []byte { | |||
| _, _ = task.UpdateWithStatus(snap.Status) | |||
| } | |||
| // OpenAI Video API 由调用者的 ConvertToOpenAIVideo 分支处理 | |||
| if isOpenAIVideoAPI { | |||
| return nil | |||
| } | |||
| // 非 OpenAI Video API: 构建自定义格式响应 | |||
| format := detectVideoFormat(body) | |||
| out := map[string]any{ | |||
| "error": nil, | |||
| @@ -493,7 +463,6 @@ func tryRealtimeFetch(task *model.Task, isOpenAIVideoAPI bool) []byte { | |||
| return respBody | |||
| } | |||
| // detectVideoFormat 从 Gemini/Vertex 原始响应中探测视频格式 | |||
| func detectVideoFormat(rawBody []byte) string { | |||
| var raw map[string]any | |||
| if err := common.Unmarshal(rawBody, &raw); err != nil { | |||
| @@ -518,7 +487,6 @@ func detectVideoFormat(rawBody []byte) string { | |||
| return mt | |||
| } | |||
| // mapTaskStatusToSimple 将内部 TaskStatus 映射为简化状态字符串 | |||
| func mapTaskStatusToSimple(status model.TaskStatus) string { | |||
| switch status { | |||
| case model.TaskStatusSuccess: | |||
| @@ -553,6 +521,26 @@ func TaskModel2Dto(task *model.Task) *dto.TaskDto { | |||
| Progress: task.Progress, | |||
| Properties: task.Properties, | |||
| Username: task.Username, | |||
| Data: task.Data, | |||
| Data: sanitizeTaskDtoData(task), | |||
| } | |||
| } | |||
| func sanitizeTaskDtoData(task *model.Task) []byte { | |||
| if task == nil || len(task.Data) == 0 { | |||
| return nil | |||
| } | |||
| payload := map[string]any{} | |||
| if err := common.Unmarshal(task.Data, &payload); err != nil { | |||
| return task.Data | |||
| } | |||
| if _, ok := payload["aiping_id"]; !ok { | |||
| return task.Data | |||
| } | |||
| delete(payload, "aiping_id") | |||
| payload["id"] = task.TaskID | |||
| data, err := common.Marshal(payload) | |||
| if err != nil { | |||
| return task.Data | |||
| } | |||
| return data | |||
| } | |||
| @@ -0,0 +1,230 @@ | |||
| package relay | |||
| import ( | |||
| "net/http" | |||
| "net/http/httptest" | |||
| "strings" | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/constant" | |||
| "github.com/QuantumNous/new-api/model" | |||
| relaycommon "github.com/QuantumNous/new-api/relay/common" | |||
| "github.com/QuantumNous/new-api/types" | |||
| "github.com/gin-gonic/gin" | |||
| "github.com/glebarez/sqlite" | |||
| "github.com/stretchr/testify/assert" | |||
| "github.com/stretchr/testify/require" | |||
| "gorm.io/gorm" | |||
| ) | |||
| func setupRelayTaskTestDB(t *testing.T) { | |||
| t.Helper() | |||
| db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) | |||
| require.NoError(t, err) | |||
| require.NoError(t, db.AutoMigrate(&model.Task{}, &model.Channel{})) | |||
| sqlDB, err := db.DB() | |||
| require.NoError(t, err) | |||
| sqlDB.SetMaxOpenConns(1) | |||
| origDB := model.DB | |||
| origUsingSQLite := common.UsingSQLite | |||
| origRedisEnabled := common.RedisEnabled | |||
| model.DB = db | |||
| common.UsingSQLite = true | |||
| common.RedisEnabled = false | |||
| t.Cleanup(func() { | |||
| model.DB = origDB | |||
| common.UsingSQLite = origUsingSQLite | |||
| common.RedisEnabled = origRedisEnabled | |||
| require.NoError(t, sqlDB.Close()) | |||
| }) | |||
| } | |||
| func TestResolveOriginTask_RestoresMatrixOriginPricing(t *testing.T) { | |||
| setupRelayTaskTestDB(t) | |||
| require.NoError(t, model.DB.Create(&model.Channel{ | |||
| Id: 11, | |||
| Type: constant.ChannelTypeMiniMax, | |||
| Key: "test-key", | |||
| Status: common.ChannelStatusEnabled, | |||
| Name: "hailuo", | |||
| Group: "default", | |||
| Models: "matrix-video", | |||
| }).Error) | |||
| require.NoError(t, model.DB.Create(&model.Task{ | |||
| TaskID: "task_origin", | |||
| UserId: 7, | |||
| ChannelId: 11, | |||
| Properties: model.Properties{ | |||
| OriginModelName: "matrix-video", | |||
| }, | |||
| PrivateData: model.TaskPrivateData{ | |||
| BillingContext: &model.TaskBillingContext{ | |||
| BillingMode: types.BillingModeMatrix, | |||
| BillingUnit: types.BillingUnitPer1MTokens, | |||
| TokenUnitPriceUSD: 0.5, | |||
| PricingSnapshot: map[string]any{ | |||
| "resolution": "720P", | |||
| "price": 0.5, | |||
| "billing_unit": types.BillingUnitPer1MTokens, | |||
| "pricing_mode": string(types.PricingModeMatched), | |||
| }, | |||
| PerCallBilling: false, | |||
| }, | |||
| }, | |||
| }).Error) | |||
| c := newRelayTaskRemixContext(t, "task_origin") | |||
| info := &relaycommon.RelayInfo{UserId: 7, TaskRelayInfo: &relaycommon.TaskRelayInfo{}} | |||
| taskErr := ResolveOriginTask(c, info) | |||
| require.Nil(t, taskErr) | |||
| require.NotNil(t, info.OriginPricing) | |||
| require.NotNil(t, info.LockedChannel) | |||
| assert.Equal(t, constant.TaskActionRemix, info.Action) | |||
| assert.Equal(t, "matrix-video", info.OriginModelName) | |||
| assert.Equal(t, types.BillingModeMatrix, info.OriginPricing.BillingMode) | |||
| assert.Equal(t, types.BillingUnitPer1MTokens, info.OriginPricing.BillingUnit) | |||
| assert.Equal(t, 0.5, info.OriginPricing.TokenUnitPriceUSD) | |||
| assert.Equal(t, "720P", info.OriginPricing.Snapshot["resolution"]) | |||
| lockedChannel, ok := info.LockedChannel.(*model.Channel) | |||
| require.True(t, ok) | |||
| assert.Equal(t, 11, lockedChannel.Id) | |||
| } | |||
| func TestResolveOriginTask_RestoresLegacyOriginPricing(t *testing.T) { | |||
| setupRelayTaskTestDB(t) | |||
| require.NoError(t, model.DB.Create(&model.Channel{ | |||
| Id: 12, | |||
| Type: constant.ChannelTypeOpenAI, | |||
| Key: "test-key", | |||
| Status: common.ChannelStatusEnabled, | |||
| Name: "sora", | |||
| Group: "default", | |||
| Models: "legacy-video", | |||
| }).Error) | |||
| require.NoError(t, model.DB.Create(&model.Task{ | |||
| TaskID: "task_legacy", | |||
| UserId: 7, | |||
| ChannelId: 12, | |||
| Properties: model.Properties{ | |||
| OriginModelName: "legacy-video", | |||
| }, | |||
| PrivateData: model.TaskPrivateData{ | |||
| BillingContext: &model.TaskBillingContext{ | |||
| ModelPrice: 0.2, | |||
| OtherRatios: map[string]float64{"seconds": 10, "size": 1.666667}, | |||
| PerCallBilling: true, | |||
| }, | |||
| }, | |||
| }).Error) | |||
| c := newRelayTaskRemixContext(t, "task_legacy") | |||
| info := &relaycommon.RelayInfo{UserId: 7, TaskRelayInfo: &relaycommon.TaskRelayInfo{}} | |||
| taskErr := ResolveOriginTask(c, info) | |||
| require.Nil(t, taskErr) | |||
| require.NotNil(t, info.OriginPricing) | |||
| require.NotNil(t, info.LockedChannel) | |||
| assert.Equal(t, "legacy-video", info.OriginModelName) | |||
| assert.Equal(t, 0.2, info.OriginPricing.PriceUSD) | |||
| assert.Equal(t, map[string]float64{"seconds": 10, "size": 1.666667}, info.OriginPricing.OtherRatios) | |||
| assert.True(t, info.OriginPricing.PerCallBilling) | |||
| lockedChannel, ok := info.LockedChannel.(*model.Channel) | |||
| require.True(t, ok) | |||
| assert.Equal(t, 12, lockedChannel.Id) | |||
| } | |||
| func TestBuildUpstreamRequestSnapshotRedactsDataURIsAndKeepsDebugFields(t *testing.T) { | |||
| raw := []byte(`{ | |||
| "model":"doubao-seedance-2-0-260128", | |||
| "content":[ | |||
| {"type":"text","text":"prompt"}, | |||
| {"type":"image_url","image_url":{"url":"data:image/png;base64,QUJDREVGRw=="}}, | |||
| {"type":"video_url","video_url":{"url":"asset://video-1"}}, | |||
| {"type":"audio_url","audio_url":{"url":"https://example.test/audio.wav"}} | |||
| ] | |||
| }`) | |||
| snapshot := buildUpstreamRequestSnapshot(raw) | |||
| require.NotNil(t, snapshot) | |||
| require.Equal(t, len(raw), snapshot.RawBytes) | |||
| require.NotEmpty(t, snapshot.SHA256) | |||
| require.True(t, snapshot.Redacted) | |||
| require.False(t, snapshot.Truncated) | |||
| data, err := common.Marshal(snapshot) | |||
| require.NoError(t, err) | |||
| jsonBody := string(data) | |||
| require.Contains(t, jsonBody, `"model":"doubao-seedance-2-0-260128"`) | |||
| require.Contains(t, jsonBody, `"asset://video-1"`) | |||
| require.Contains(t, jsonBody, `"https://example.test/audio.wav"`) | |||
| require.Contains(t, jsonBody, `"kind":"data_uri"`) | |||
| require.Contains(t, jsonBody, `"media_type":"image/png"`) | |||
| require.NotContains(t, jsonBody, "QUJDREVGRw==") | |||
| require.Equal(t, len(data), snapshot.StoredBytes) | |||
| } | |||
| func TestBuildUpstreamRequestSnapshotTruncatesOversizedSnapshot(t *testing.T) { | |||
| raw := []byte(`{"model":"m","prompt":"` + strings.Repeat("x", maxUpstreamRequestSnapshotBytes+1024) + `"}`) | |||
| snapshot := buildUpstreamRequestSnapshot(raw) | |||
| require.NotNil(t, snapshot) | |||
| require.True(t, snapshot.Truncated) | |||
| require.NotEmpty(t, snapshot.SHA256) | |||
| data, err := common.Marshal(snapshot) | |||
| require.NoError(t, err) | |||
| require.LessOrEqual(t, len(data), maxUpstreamRequestSnapshotBytes) | |||
| require.Equal(t, len(data), snapshot.StoredBytes) | |||
| } | |||
| func TestTaskModel2DtoDoesNotExposeUpstreamRequestSnapshot(t *testing.T) { | |||
| task := &model.Task{ | |||
| TaskID: "task_public", | |||
| PrivateData: model.TaskPrivateData{ | |||
| UpstreamRequest: &model.TaskUpstreamRequestSnapshot{ | |||
| Body: map[string]any{"secret": "hidden"}, | |||
| }, | |||
| }, | |||
| } | |||
| dtoTask := TaskModel2Dto(task) | |||
| data, err := common.Marshal(dtoTask) | |||
| require.NoError(t, err) | |||
| require.NotContains(t, string(data), "upstream_request") | |||
| require.NotContains(t, string(data), "hidden") | |||
| } | |||
| func TestTaskModel2DtoSanitizesAipingInternalIDFromData(t *testing.T) { | |||
| task := &model.Task{ | |||
| TaskID: "task_public", | |||
| Data: []byte(`{ | |||
| "id":"cgt-upstream", | |||
| "aiping_id":"b7ee658f-acac-433d-897d-5c9ed55ef402", | |||
| "status":"succeeded" | |||
| }`), | |||
| } | |||
| dtoTask := TaskModel2Dto(task) | |||
| data, err := common.Marshal(dtoTask) | |||
| require.NoError(t, err) | |||
| require.NotContains(t, string(data), "aiping_id") | |||
| require.NotContains(t, string(data), "b7ee658f-acac-433d-897d-5c9ed55ef402") | |||
| require.Contains(t, string(data), `"id":"task_public"`) | |||
| } | |||
| func newRelayTaskRemixContext(t *testing.T, originTaskID string) *gin.Context { | |||
| t.Helper() | |||
| gin.SetMode(gin.TestMode) | |||
| w := httptest.NewRecorder() | |||
| c, _ := gin.CreateTestContext(w) | |||
| c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos/"+originTaskID+"/remix", nil) | |||
| c.Params = gin.Params{{Key: "video_id", Value: originTaskID}} | |||
| return c | |||
| } | |||
| @@ -0,0 +1,150 @@ | |||
| package relay | |||
| import ( | |||
| "crypto/sha256" | |||
| "encoding/hex" | |||
| "strings" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/model" | |||
| ) | |||
| const maxUpstreamRequestSnapshotBytes = 256 * 1024 | |||
| func buildUpstreamRequestSnapshot(raw []byte) *model.TaskUpstreamRequestSnapshot { | |||
| if len(raw) == 0 { | |||
| return nil | |||
| } | |||
| snapshot := &model.TaskUpstreamRequestSnapshot{ | |||
| RawBytes: len(raw), | |||
| SHA256: sha256Hex(raw), | |||
| } | |||
| var body any | |||
| if err := common.Unmarshal(raw, &body); err == nil { | |||
| sanitized, redacted := sanitizeUpstreamRequestValue(body) | |||
| snapshot.Body = sanitized | |||
| snapshot.Redacted = redacted | |||
| } else { | |||
| snapshot.Body = map[string]any{ | |||
| "raw_preview": truncateStringToBytes(string(raw), maxUpstreamRequestSnapshotBytes/2), | |||
| "format": "non_json", | |||
| } | |||
| snapshot.Truncated = len(raw) > maxUpstreamRequestSnapshotBytes/2 | |||
| } | |||
| finalizeUpstreamRequestSnapshot(snapshot, raw) | |||
| return snapshot | |||
| } | |||
| func BuildUpstreamRequestSnapshotForTask(raw []byte) *model.TaskUpstreamRequestSnapshot { | |||
| return buildUpstreamRequestSnapshot(raw) | |||
| } | |||
| func sanitizeUpstreamRequestValue(value any) (any, bool) { | |||
| switch v := value.(type) { | |||
| case map[string]any: | |||
| out := make(map[string]any, len(v)) | |||
| redacted := false | |||
| for key, item := range v { | |||
| sanitized, itemRedacted := sanitizeUpstreamRequestValue(item) | |||
| out[key] = sanitized | |||
| redacted = redacted || itemRedacted | |||
| } | |||
| return out, redacted | |||
| case []any: | |||
| out := make([]any, len(v)) | |||
| redacted := false | |||
| for i, item := range v { | |||
| sanitized, itemRedacted := sanitizeUpstreamRequestValue(item) | |||
| out[i] = sanitized | |||
| redacted = redacted || itemRedacted | |||
| } | |||
| return out, redacted | |||
| case string: | |||
| if mediaType, ok := parseBase64DataURI(v); ok { | |||
| return map[string]any{ | |||
| "redacted": true, | |||
| "kind": "data_uri", | |||
| "media_type": mediaType, | |||
| "bytes": len(v), | |||
| "sha256": sha256Hex([]byte(v)), | |||
| }, true | |||
| } | |||
| return v, false | |||
| default: | |||
| return value, false | |||
| } | |||
| } | |||
| func parseBase64DataURI(value string) (string, bool) { | |||
| lower := strings.ToLower(value) | |||
| if !strings.HasPrefix(lower, "data:") { | |||
| return "", false | |||
| } | |||
| marker := ";base64," | |||
| idx := strings.Index(lower, marker) | |||
| if idx <= len("data:") { | |||
| return "", false | |||
| } | |||
| mediaType := value[len("data:"):idx] | |||
| if !strings.HasPrefix(strings.ToLower(mediaType), "image/") && | |||
| !strings.HasPrefix(strings.ToLower(mediaType), "audio/") && | |||
| !strings.HasPrefix(strings.ToLower(mediaType), "video/") { | |||
| return "", false | |||
| } | |||
| return mediaType, true | |||
| } | |||
| func finalizeUpstreamRequestSnapshot(snapshot *model.TaskUpstreamRequestSnapshot, raw []byte) { | |||
| updateStoredBytes(snapshot) | |||
| if snapshot.StoredBytes <= maxUpstreamRequestSnapshotBytes { | |||
| return | |||
| } | |||
| bodyBytes, err := common.Marshal(snapshot.Body) | |||
| previewSource := raw | |||
| if err == nil { | |||
| previewSource = bodyBytes | |||
| } | |||
| snapshot.Body = map[string]any{ | |||
| "truncated_preview": truncateStringToBytes(string(previewSource), maxUpstreamRequestSnapshotBytes/2), | |||
| } | |||
| snapshot.Truncated = true | |||
| updateStoredBytes(snapshot) | |||
| for snapshot.StoredBytes > maxUpstreamRequestSnapshotBytes { | |||
| body, _ := snapshot.Body.(map[string]any) | |||
| preview, _ := body["truncated_preview"].(string) | |||
| if len(preview) == 0 { | |||
| break | |||
| } | |||
| body["truncated_preview"] = truncateStringToBytes(preview, len(preview)/2) | |||
| snapshot.Body = body | |||
| updateStoredBytes(snapshot) | |||
| } | |||
| } | |||
| func updateStoredBytes(snapshot *model.TaskUpstreamRequestSnapshot) { | |||
| for i := 0; i < 4; i++ { | |||
| data, err := common.Marshal(snapshot) | |||
| if err != nil { | |||
| return | |||
| } | |||
| if snapshot.StoredBytes == len(data) { | |||
| return | |||
| } | |||
| snapshot.StoredBytes = len(data) | |||
| } | |||
| } | |||
| func truncateStringToBytes(value string, maxBytes int) string { | |||
| if maxBytes <= 0 || len(value) <= maxBytes { | |||
| return value | |||
| } | |||
| return value[:maxBytes] | |||
| } | |||
| func sha256Hex(data []byte) string { | |||
| sum := sha256.Sum256(data) | |||
| return hex.EncodeToString(sum[:]) | |||
| } | |||
| @@ -50,6 +50,12 @@ func SetApiRouter(router *gin.Engine) { | |||
| apiRouter.GET("/ratio_config", middleware.CriticalRateLimit(), controller.GetRatioConfig) | |||
| apiRouter.GET("/playground/config", controller.GetPlaygroundConfig) | |||
| assetRoute := apiRouter.Group("/v1/volcengine") | |||
| assetRoute.Use(middleware.SystemPerformanceCheck(), middleware.TokenAuth()) | |||
| { | |||
| assetRoute.POST("/asset", controller.DoubaoAssetProxy) | |||
| } | |||
| apiRouter.POST("/stripe/webhook", controller.StripeWebhook) | |||
| apiRouter.POST("/creem/webhook", controller.CreemWebhook) | |||
| apiRouter.POST("/wechat/pay/webhook", controller.WechatPayWebhook) | |||
| @@ -185,6 +191,14 @@ func SetApiRouter(router *gin.Engine) { | |||
| { | |||
| optionRoute.GET("/", controller.GetOptions) | |||
| optionRoute.PUT("/", controller.UpdateOption) | |||
| optionRoute.GET("/model_pricing", controller.GetModelPricingRules) | |||
| optionRoute.GET("/model_pricing/*model", controller.GetModelPricingRule) | |||
| optionRoute.PUT("/model_pricing/*model", controller.UpdateModelPricingRule) | |||
| optionRoute.DELETE("/model_pricing/*model", controller.DeleteModelPricingRule) | |||
| optionRoute.GET("/model_display_pricing", controller.GetModelDisplayPricingRules) | |||
| optionRoute.GET("/model_display_pricing/*model", controller.GetModelDisplayPricingRule) | |||
| optionRoute.PUT("/model_display_pricing/*model", controller.UpdateModelDisplayPricingRule) | |||
| optionRoute.DELETE("/model_display_pricing/*model", controller.DeleteModelDisplayPricingRule) | |||
| optionRoute.GET("/channel_affinity_cache", controller.GetChannelAffinityCacheStats) | |||
| optionRoute.DELETE("/channel_affinity_cache", controller.ClearChannelAffinityCache) | |||
| optionRoute.POST("/rest_model_ratio", controller.ResetModelRatio) | |||
| @@ -0,0 +1,188 @@ | |||
| package router | |||
| import ( | |||
| "net/http" | |||
| "net/http/httptest" | |||
| "os" | |||
| "strings" | |||
| "sync/atomic" | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/constant" | |||
| "github.com/QuantumNous/new-api/model" | |||
| "github.com/QuantumNous/new-api/service" | |||
| "github.com/QuantumNous/new-api/setting/ratio_setting" | |||
| "github.com/gin-gonic/gin" | |||
| "github.com/stretchr/testify/require" | |||
| "gorm.io/gorm" | |||
| "gorm.io/gorm/logger" | |||
| ) | |||
| func setupTianyiYunSeedanceE2EDB(t *testing.T) *gorm.DB { | |||
| t.Helper() | |||
| oldDB := model.DB | |||
| oldLOGDB := model.LOG_DB | |||
| oldSQLitePath := common.SQLitePath | |||
| oldMemoryCacheEnabled := common.MemoryCacheEnabled | |||
| oldRedisEnabled := common.RedisEnabled | |||
| oldIsMasterNode := common.IsMasterNode | |||
| oldUsingSQLite := common.UsingSQLite | |||
| oldUsingMySQL := common.UsingMySQL | |||
| oldUsingPostgreSQL := common.UsingPostgreSQL | |||
| oldSQLDSN, hadSQLDSN := os.LookupEnv("SQL_DSN") | |||
| oldModelRatio := ratio_setting.ModelRatio2JSONString() | |||
| oldGroupRatio := ratio_setting.GroupRatio2JSONString() | |||
| common.SQLitePath = "file:" + strings.NewReplacer("/", "_", " ", "_").Replace(t.Name()) + "?mode=memory&cache=shared" | |||
| common.MemoryCacheEnabled = false | |||
| common.RedisEnabled = false | |||
| common.IsMasterNode = false | |||
| common.UsingSQLite = false | |||
| common.UsingMySQL = false | |||
| common.UsingPostgreSQL = false | |||
| require.NoError(t, os.Setenv("SQL_DSN", "local")) | |||
| require.NoError(t, model.InitDB()) | |||
| require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"cdance2.0-0611":0}`)) | |||
| require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1}`)) | |||
| service.InitHttpClient() | |||
| model.DB = model.DB.Session(&gorm.Session{Logger: logger.Default.LogMode(logger.Silent)}) | |||
| model.LOG_DB = model.DB | |||
| db := model.DB | |||
| sqlDB, err := db.DB() | |||
| require.NoError(t, err) | |||
| sqlDB.SetMaxOpenConns(1) | |||
| require.NoError(t, db.AutoMigrate( | |||
| &model.User{}, | |||
| &model.Token{}, | |||
| &model.Channel{}, | |||
| &model.Ability{}, | |||
| &model.UserAssetChannel{}, | |||
| &model.Task{}, | |||
| &model.Log{}, | |||
| &model.SubscriptionPlan{}, | |||
| &model.SubscriptionOrder{}, | |||
| &model.UserSubscription{}, | |||
| &model.SubscriptionPreConsumeRecord{}, | |||
| )) | |||
| t.Cleanup(func() { | |||
| _ = sqlDB.Close() | |||
| model.DB = oldDB | |||
| model.LOG_DB = oldLOGDB | |||
| common.SQLitePath = oldSQLitePath | |||
| common.MemoryCacheEnabled = oldMemoryCacheEnabled | |||
| common.RedisEnabled = oldRedisEnabled | |||
| common.IsMasterNode = oldIsMasterNode | |||
| common.UsingSQLite = oldUsingSQLite | |||
| common.UsingMySQL = oldUsingMySQL | |||
| common.UsingPostgreSQL = oldUsingPostgreSQL | |||
| _ = ratio_setting.UpdateModelRatioByJSONString(oldModelRatio) | |||
| _ = ratio_setting.UpdateGroupRatioByJSONString(oldGroupRatio) | |||
| if hadSQLDSN { | |||
| _ = os.Setenv("SQL_DSN", oldSQLDSN) | |||
| } else { | |||
| _ = os.Unsetenv("SQL_DSN") | |||
| } | |||
| }) | |||
| return db | |||
| } | |||
| func TestTianyiYunSeedanceSubmitE2EUsesBoundUserChannel(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| db := setupTianyiYunSeedanceE2EDB(t) | |||
| var upstreamCalls int32 | |||
| var gotPath string | |||
| var gotAuth string | |||
| upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | |||
| atomic.AddInt32(&upstreamCalls, 1) | |||
| gotPath = r.URL.Path | |||
| gotAuth = r.Header.Get("Authorization") | |||
| w.Header().Set("Content-Type", "application/json") | |||
| _, _ = w.Write([]byte(`{"id":"task-upstream","status":"queued","model":"cdance2.0-0611"}`)) | |||
| })) | |||
| t.Cleanup(upstream.Close) | |||
| require.NoError(t, db.Create(&model.User{ | |||
| Id: 10, | |||
| Username: "e2e-user", | |||
| Status: common.UserStatusEnabled, | |||
| Group: "default", | |||
| Quota: 1000000, | |||
| }).Error) | |||
| require.NoError(t, db.Create(&model.Token{ | |||
| Id: 20, | |||
| UserId: 10, | |||
| Key: "e2etokenkey", | |||
| Status: common.TokenStatusEnabled, | |||
| Name: "e2e-token", | |||
| ExpiredTime: -1, | |||
| UnlimitedQuota: true, | |||
| Group: "default", | |||
| }).Error) | |||
| priority := int64(1) | |||
| weight := uint(10) | |||
| autoBan := 1 | |||
| baseURL := upstream.URL | |||
| require.NoError(t, db.Create(&model.Channel{ | |||
| Id: 30, | |||
| Type: constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, | |||
| Key: "upstream-secret", | |||
| Status: common.ChannelStatusEnabled, | |||
| Name: "tianyiyun-e2e", | |||
| Group: "default", | |||
| Models: "cdance2.0-0611", | |||
| BaseURL: &baseURL, | |||
| Priority: &priority, | |||
| Weight: &weight, | |||
| AutoBan: &autoBan, | |||
| CreatedTime: 30, | |||
| }).Error) | |||
| require.NoError(t, db.Create(&model.Ability{ | |||
| Group: "default", | |||
| Model: "cdance2.0-0611", | |||
| ChannelId: 30, | |||
| Enabled: true, | |||
| Priority: &priority, | |||
| Weight: 10, | |||
| }).Error) | |||
| require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default", 30)) | |||
| r := gin.New() | |||
| SetVideoRouter(r) | |||
| body := `{ | |||
| "model":"cdance2.0-0611", | |||
| "content":[{"type":"text","text":"e2e prompt"}], | |||
| "ratio":"16:9", | |||
| "duration":5, | |||
| "watermark":false | |||
| }` | |||
| req := httptest.NewRequest(http.MethodPost, "/api/v3/contents/generations/tasks", strings.NewReader(body)) | |||
| req.Header.Set("Content-Type", "application/json") | |||
| req.Header.Set("Authorization", "Bearer sk-e2etokenkey") | |||
| w := httptest.NewRecorder() | |||
| r.ServeHTTP(w, req) | |||
| require.Equal(t, http.StatusOK, w.Code, w.Body.String()) | |||
| require.Equal(t, int32(1), atomic.LoadInt32(&upstreamCalls)) | |||
| require.Equal(t, "/v1/contents/generations/tasks", gotPath) | |||
| require.Equal(t, "Bearer upstream-secret", gotAuth) | |||
| require.Contains(t, w.Body.String(), `"id":"task_`) | |||
| binding, err := model.GetUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default") | |||
| require.NoError(t, err) | |||
| require.NotNil(t, binding) | |||
| require.Equal(t, 30, binding.ChannelId) | |||
| var task model.Task | |||
| require.NoError(t, db.Where("user_id = ? AND channel_id = ?", 10, 30).First(&task).Error) | |||
| require.Equal(t, constant.TaskPlatform("60"), task.Platform) | |||
| require.Equal(t, "task-upstream", task.PrivateData.UpstreamTaskID) | |||
| require.Equal(t, "cdance2.0-0611", task.Properties.OriginModelName) | |||
| } | |||
| @@ -22,6 +22,40 @@ func SetVideoRouter(router *gin.Engine) { | |||
| videoV1Router.GET("/video/generations/:task_id", controller.RelayTaskFetch) | |||
| videoV1Router.POST("/videos/:video_id/remix", controller.RelayTask) | |||
| } | |||
| klingAipingRouter := router.Group("/v1") | |||
| klingAipingRouter.Use(middleware.TokenAuth()) | |||
| { | |||
| klingAipingRouter.POST("/videos/text2video", controller.KlingAipingNativeTaskSubmit) | |||
| klingAipingRouter.GET("/videos/text2video/:task_id", controller.KlingAipingNativeTaskFetch) | |||
| klingAipingRouter.GET("/videos/text2video", controller.KlingAipingNativeTaskList) | |||
| klingAipingRouter.POST("/videos/image2video", controller.KlingAipingNativeTaskSubmit) | |||
| klingAipingRouter.GET("/videos/image2video/:task_id", controller.KlingAipingNativeTaskFetch) | |||
| klingAipingRouter.GET("/videos/image2video", controller.KlingAipingNativeTaskList) | |||
| klingAipingRouter.POST("/videos/motion-control", controller.KlingAipingNativeTaskSubmit) | |||
| klingAipingRouter.GET("/videos/motion-control/:task_id", controller.KlingAipingNativeTaskFetch) | |||
| klingAipingRouter.GET("/videos/motion-control", controller.KlingAipingNativeTaskList) | |||
| klingAipingRouter.POST("/videos/omni-video", controller.KlingAipingNativeTaskSubmit) | |||
| klingAipingRouter.GET("/videos/omni-video/:task_id", controller.KlingAipingNativeTaskFetch) | |||
| klingAipingRouter.GET("/videos/omni-video", controller.KlingAipingNativeTaskList) | |||
| klingAipingRouter.POST("/videos/multi-image2video", controller.KlingAipingNativeTaskSubmit) | |||
| klingAipingRouter.GET("/videos/multi-image2video/:task_id", controller.KlingAipingNativeTaskFetch) | |||
| klingAipingRouter.GET("/videos/multi-image2video", controller.KlingAipingNativeTaskList) | |||
| klingAipingRouter.POST("/videos/video-extend", controller.KlingAipingNativeTaskSubmit) | |||
| klingAipingRouter.GET("/videos/video-extend/:task_id", controller.KlingAipingNativeTaskFetch) | |||
| klingAipingRouter.GET("/videos/video-extend", controller.KlingAipingNativeTaskList) | |||
| klingAipingRouter.POST("/general/advanced-custom-elements", controller.KlingAipingNativeTaskSubmit) | |||
| klingAipingRouter.GET("/general/advanced-custom-elements/:task_id", controller.KlingAipingNativeTaskFetch) | |||
| klingAipingRouter.GET("/general/advanced-custom-elements", controller.KlingAipingNativeProxy) | |||
| klingAipingRouter.GET("/general/advanced-presets-elements", controller.KlingAipingNativeProxy) | |||
| klingAipingRouter.POST("/general/delete-advanced-elements", controller.KlingAipingNativeProxy) | |||
| klingAipingRouter.POST("/general/custom-voices", controller.KlingAipingNativeTaskSubmit) | |||
| klingAipingRouter.GET("/general/custom-voices/:task_id", controller.KlingAipingNativeTaskFetch) | |||
| klingAipingRouter.GET("/general/custom-voices", controller.KlingAipingNativeProxy) | |||
| klingAipingRouter.GET("/general/presets-voices", controller.KlingAipingNativeProxy) | |||
| klingAipingRouter.POST("/general/delete-voices", controller.KlingAipingNativeProxy) | |||
| } | |||
| // openai compatible API video routes | |||
| // docs: https://platform.openai.com/docs/api-reference/videos/create | |||
| { | |||
| @@ -29,6 +63,18 @@ func SetVideoRouter(router *gin.Engine) { | |||
| videoV1Router.GET("/videos/:task_id", controller.RelayTaskFetch) | |||
| } | |||
| aipingNativeRouter := router.Group("/api/v3/contents/generations") | |||
| aipingNativeRouter.Use(middleware.TokenAuth(), middleware.Distribute()) | |||
| { | |||
| aipingNativeRouter.POST("/tasks", controller.AipingNativeVideoSubmit) | |||
| } | |||
| // Fetch 不需要 Distribute(无请求体分发逻辑) | |||
| aipingNativeFetchRouter := router.Group("/api/v3/contents/generations") | |||
| aipingNativeFetchRouter.Use(middleware.TokenAuth()) | |||
| { | |||
| aipingNativeFetchRouter.GET("/tasks/:task_id", controller.AipingNativeVideoFetch) | |||
| } | |||
| klingV1Router := router.Group("/kling/v1") | |||
| klingV1Router.Use(middleware.KlingRequestConvert(), middleware.TokenAuth(), middleware.Distribute()) | |||
| { | |||
| @@ -0,0 +1,41 @@ | |||
| package router | |||
| import ( | |||
| "net/http" | |||
| "strings" | |||
| "testing" | |||
| "github.com/gin-gonic/gin" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestKlingAipingNativeRoutesDoNotFallThroughToOpenAIVideoRoutes(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| r := gin.New() | |||
| SetVideoRouter(r) | |||
| requireHandlerContains(t, r, http.MethodPost, "/v1/videos/text2video", "KlingAipingNativeTaskSubmit") | |||
| requireHandlerContains(t, r, http.MethodGet, "/v1/videos/text2video/:task_id", "KlingAipingNativeTaskFetch") | |||
| requireHandlerContains(t, r, http.MethodGet, "/v1/videos/text2video", "KlingAipingNativeTaskList") | |||
| requireHandlerContains(t, r, http.MethodPost, "/v1/videos/video-extend", "KlingAipingNativeTaskSubmit") | |||
| requireHandlerContains(t, r, http.MethodGet, "/v1/videos/video-extend/:task_id", "KlingAipingNativeTaskFetch") | |||
| requireHandlerContains(t, r, http.MethodGet, "/v1/videos/video-extend", "KlingAipingNativeTaskList") | |||
| requireHandlerContains(t, r, http.MethodGet, "/v1/general/advanced-presets-elements", "KlingAipingNativeProxy") | |||
| requireHandlerContains(t, r, http.MethodPost, "/v1/general/delete-advanced-elements", "KlingAipingNativeProxy") | |||
| requireHandlerContains(t, r, http.MethodPost, "/v1/general/delete-voices", "KlingAipingNativeProxy") | |||
| requireHandlerContains(t, r, http.MethodPost, "/v1/videos", "RelayTask") | |||
| requireHandlerContains(t, r, http.MethodGet, "/v1/videos/:task_id", "RelayTaskFetch") | |||
| } | |||
| func requireHandlerContains(t *testing.T, r *gin.Engine, method, path, want string) { | |||
| t.Helper() | |||
| for _, route := range r.Routes() { | |||
| if route.Method != method || route.Path != path { | |||
| continue | |||
| } | |||
| require.Truef(t, strings.Contains(route.Handler, want), "handler for %s %s = %s, want containing %s", method, path, route.Handler, want) | |||
| return | |||
| } | |||
| t.Fatalf("route not found: %s %s", method, path) | |||
| } | |||