From 86470628df5c30d6cec5ee176595b0d43a2c505a Mon Sep 17 00:00:00 2001 From: alfadb Date: Sat, 22 Aug 2026 12:38:05 +0800 Subject: [PATCH] =?UTF-8?q?feat(ollama):=20=E5=AF=B9=20Ollama=20Cloud=20?= =?UTF-8?q?=E8=B4=A6=E5=8F=B7=20clamp=20max=5Ftokens=20=E4=B8=8A=E9=99=90?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- ...penai_gateway_ollama_cloud_cc_reasoning.go | 3 +- .../openai_gateway_ollama_cloud_max_tokens.go | 83 +++++++++ ...ai_gateway_ollama_cloud_max_tokens_test.go | 160 ++++++++++++++++++ 3 files changed, 245 insertions(+), 1 deletion(-) create mode 100644 backend/internal/service/openai_gateway_ollama_cloud_max_tokens.go create mode 100644 backend/internal/service/openai_gateway_ollama_cloud_max_tokens_test.go diff --git a/backend/internal/service/openai_gateway_ollama_cloud_cc_reasoning.go b/backend/internal/service/openai_gateway_ollama_cloud_cc_reasoning.go index 659ea4f5d2..f78a375860 100644 --- a/backend/internal/service/openai_gateway_ollama_cloud_cc_reasoning.go +++ b/backend/internal/service/openai_gateway_ollama_cloud_cc_reasoning.go @@ -51,7 +51,8 @@ func applyOllamaCloudRawChatCompletionsRequest(account *Account, body []byte) [] if !isOllamaCloudRawChatCompletionsAccount(account) || len(body) == 0 { return body } - return normalizeOllamaCloudChatCompletionsRequest(body) + body = normalizeOllamaCloudChatCompletionsRequest(body) + return clampOllamaCloudMaxTokens(account, body) } func applyOllamaCloudRawChatCompletionsResponse(account *Account, body []byte) []byte { diff --git a/backend/internal/service/openai_gateway_ollama_cloud_max_tokens.go b/backend/internal/service/openai_gateway_ollama_cloud_max_tokens.go new file mode 100644 index 0000000000..32a508481a --- /dev/null +++ b/backend/internal/service/openai_gateway_ollama_cloud_max_tokens.go @@ -0,0 +1,83 @@ +package service + +import ( + "encoding/json" + + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" + "go.uber.org/zap" +) + +// OllamaCloudMaxTokensCapExtraKey 是账号 extra 中的可选配置键，表示该 Ollama Cloud +// 账号输出 token 的 provider 级硬上限。用户可通过 admin 账号更新 API 的 extra 字段 +// 设置，覆盖默认值 ollamaCloudDefaultMaxTokensCap；0 或负数表示显式禁用 clamp。 +const OllamaCloudMaxTokensCapExtraKey = "ollama_max_tokens_cap" + +// ollamaCloudDefaultMaxTokensCap 是 Ollama Cloud 对输出 token 数的 provider 级硬上限 +// （约 65535），max_tokens 超过该值会被上游直接 400 拒绝；该上限与模型无关，不做模型过滤。 +const ollamaCloudDefaultMaxTokensCap = 65535 + +// 本文件的 clampOllamaCloudMaxTokens 被 +// applyOllamaCloudRawChatCompletionsRequest（openai_gateway_ollama_cloud_cc_reasoning.go） +// 调用，账号检测（isOllamaCloudRawChatCompletionsAccount）由调用方完成，此处不再重复判断。 + +// ollamaCloudMaxTokensCap 返回账号配置的 max_tokens 上限。账号为 nil 或 extra 中 +// 无该键时返回默认值；键值为数值类型（float64/int64/int/json.Number）时返回其整数 +// 值（0 或负数表示显式禁用 clamp）；其它类型回退默认值。 +func ollamaCloudMaxTokensCap(account *Account) int64 { + if account == nil || account.Extra == nil { + return ollamaCloudDefaultMaxTokensCap + } + value, ok := account.Extra[OllamaCloudMaxTokensCapExtraKey] + if !ok { + return ollamaCloudDefaultMaxTokensCap + } + switch number := value.(type) { + case float64: + return int64(number) + case int64: + return number + case int: + return int64(number) + case json.Number: + parsed, err := number.Int64() + if err != nil { + return ollamaCloudDefaultMaxTokensCap + } + return parsed + default: + return ollamaCloudDefaultMaxTokensCap + } +} + +// clampOllamaCloudMaxTokens 把 body 中超过 cap 的 max_tokens / max_completion_tokens +// 单向压到 cap。cap <= 0 或 body 不是合法 JSON 时原样返回；sjson 出错时返回原始 body。 +// 有任一字段被 clamp 时记录一条 Debug 日志。 +func clampOllamaCloudMaxTokens(account *Account, body []byte) []byte { + cap := ollamaCloudMaxTokensCap(account) + if cap <= 0 || !gjson.ValidBytes(body) { + return body + } + clamped := false + out := body + for _, key := range []string{"max_tokens", "max_completion_tokens"} { + result := gjson.GetBytes(out, key) + if !result.Exists() || result.Type != gjson.Number || result.Int() <= cap { + continue + } + updated, err := sjson.SetBytes(out, key, cap) + if err != nil { + return body + } + out = updated + clamped = true + } + if clamped && account != nil { + logger.L().Debug("openai chat_completions raw: clamped max_tokens for ollama cloud account", + zap.Int64("account_id", account.ID), + zap.Int64("cap", cap), + ) + } + return out +} diff --git a/backend/internal/service/openai_gateway_ollama_cloud_max_tokens_test.go b/backend/internal/service/openai_gateway_ollama_cloud_max_tokens_test.go new file mode 100644 index 0000000000..6bbc2d7e5f --- /dev/null +++ b/backend/internal/service/openai_gateway_ollama_cloud_max_tokens_test.go @@ -0,0 +1,160 @@ +//go:build unit + +package service + +import ( + "encoding/json" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/pkg/openai_compat" + "github.com/stretchr/testify/require" +) + +// ollamaMaxTokensCapTestAccount 构造带自定义 cap 的 Ollama Cloud usage 账号。 +func ollamaMaxTokensCapTestAccount(id int64, cap any) *Account { + account := ollamaUsageAccount(id) + account.Extra[OllamaCloudMaxTokensCapExtraKey] = cap + return account +} + +func TestOllamaCloudMaxTokensClamp(t *testing.T) { + ollama := ollamaUsageAccount(101) + + tests := []struct { + name string + account *Account + body string + want string + raw bool // want 非法 JSON 时按原始字节比较 + }{ + { + name: "max_tokens above default cap is clamped", + account: ollama, + body: `{"model":"gpt-oss:120b-cloud","max_tokens":70000}`, + want: `{"model":"gpt-oss:120b-cloud","max_tokens":65535}`, + }, + { + name: "max_completion_tokens above default cap is clamped", + account: ollama, + body: `{"model":"gpt-oss:120b-cloud","max_completion_tokens":131072}`, + want: `{"model":"gpt-oss:120b-cloud","max_completion_tokens":65535}`, + }, + { + name: "both fields above cap are clamped", + account: ollama, + body: `{"model":"m","max_tokens":80000,"max_completion_tokens":90000}`, + want: `{"model":"m","max_tokens":65535,"max_completion_tokens":65535}`, + }, + { + name: "values at or below default cap are kept", + account: ollama, + body: `{"model":"m","max_tokens":65535,"max_completion_tokens":4096}`, + want: `{"model":"m","max_tokens":65535,"max_completion_tokens":4096}`, + }, + { + name: "custom extra cap is applied", + account: ollamaMaxTokensCapTestAccount(102, 32768), + body: `{"model":"m","max_tokens":50000}`, + want: `{"model":"m","max_tokens":32768}`, + }, + { + name: "extra cap zero disables clamping", + account: ollamaMaxTokensCapTestAccount(103, 0), + body: `{"model":"m","max_tokens":50000}`, + want: `{"model":"m","max_tokens":50000}`, + }, + { + name: "non-numeric extra cap falls back to default", + account: ollamaMaxTokensCapTestAccount(104, "abc"), + body: `{"model":"m","max_tokens":100000}`, + want: `{"model":"m","max_tokens":65535}`, + }, + { + name: "invalid json is left untouched", + account: ollama, + body: `{"model":"m","max_tokens":`, + want: `{"model":"m","max_tokens":`, + raw: true, + }, + { + name: "non-integer max_tokens is left untouched", + account: ollama, + body: `{"model":"m","max_tokens":1.5}`, + want: `{"model":"m","max_tokens":1.5}`, + }, + { + name: "missing max_tokens is left untouched", + account: ollama, + body: `{"model":"m"}`, + want: `{"model":"m"}`, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + got := clampOllamaCloudMaxTokens(test.account, []byte(test.body)) + if test.raw { + require.Equal(t, test.want, string(got)) + return + } + require.JSONEq(t, test.want, string(got)) + }) + } +} + +func TestOllamaCloudMaxTokensCap(t *testing.T) { + require.Equal(t, int64(65535), ollamaCloudMaxTokensCap(nil)) + require.Equal(t, int64(65535), ollamaCloudMaxTokensCap(ollamaUsageAccount(201))) + + tests := []struct { + name string + cap any + want int64 + }{ + {"float64", float64(32768), 32768}, + {"int", 40000, 40000}, + {"int64", int64(50000), 50000}, + {"json.Number", json.Number("60000"), 60000}, + {"json.Number invalid", json.Number("abc"), 65535}, + {"zero disables", 0, 0}, + {"negative disables", int64(-1), -1}, + {"string falls back", "abc", 65535}, + {"bool falls back", true, 65535}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + account := ollamaMaxTokensCapTestAccount(202, test.cap) + require.Equal(t, test.want, ollamaCloudMaxTokensCap(account)) + }) + } +} + +// TestApplyOllamaCloudRawChatCompletionsRequestClampsMaxTokens 验证 max_tokens clamp +// 已接入组合钩子 applyOllamaCloudRawChatCompletionsRequest，并遵循该钩子的账号判定门槛 +// （isOllamaCloudRawChatCompletionsAccount：platform openai + type apikey + +// force_chat_completions + ollama.com 或 Ollama usage extra）。 +func TestApplyOllamaCloudRawChatCompletionsRequestClampsMaxTokens(t *testing.T) { + body := []byte(`{"model":"deepseek-chat","max_tokens":100000}`) + + // Ollama Cloud 账号（ollama.com + force_chat_completions）→ clamp 到 65535。 + ollama := ollamaCloudRawChatCompletionsTestAccount() + require.JSONEq(t, `{"model":"deepseek-chat","max_tokens":65535}`, + string(applyOllamaCloudRawChatCompletionsRequest(ollama, body))) + + // 官方 DeepSeek（api.deepseek.com + force_chat_completions）→ 字节级不变。 + official := rawChatCompletionsTestAccount() + official.Credentials["base_url"] = "https://api.deepseek.com" + official.Extra = map[string]any{ + openai_compat.ExtraKeyResponsesMode: string(openai_compat.ResponsesSupportModeForceChatCompletions), + } + require.Equal(t, body, applyOllamaCloudRawChatCompletionsRequest(official, body)) + + // ollama.com 但无 force_chat_completions（Extra 缺键）→ 不通过钩子判定门槛，字节级不变。 + noForce := ollamaCloudRawChatCompletionsTestAccount() + noForce.Extra = nil + require.Equal(t, body, applyOllamaCloudRawChatCompletionsRequest(noForce, body)) + + // 空 body → 原样返回。 + require.Equal(t, []byte(nil), applyOllamaCloudRawChatCompletionsRequest(ollama, nil)) + require.Equal(t, []byte{}, applyOllamaCloudRawChatCompletionsRequest(ollama, []byte{})) +}