mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-11-07 17:57:23 +00:00
Compare commits
2 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 5b628b4d6f | |||
| 2e3d5b79ad |
4
cmd/gui/package-lock.json
generated
4
cmd/gui/package-lock.json
generated
@ -1,12 +1,12 @@
|
||||
{
|
||||
"name": "modelrouter-gui",
|
||||
"version": "1.4.1",
|
||||
"version": "1.4.2",
|
||||
"lockfileVersion": 3,
|
||||
"requires": true,
|
||||
"packages": {
|
||||
"": {
|
||||
"name": "modelrouter-gui",
|
||||
"version": "1.4.1",
|
||||
"version": "1.4.2",
|
||||
"license": "Proprietary",
|
||||
"devDependencies": {
|
||||
"electron": "^33.0.0",
|
||||
|
||||
@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "modelrouter-gui",
|
||||
"version": "1.4.1",
|
||||
"version": "1.4.2",
|
||||
"description": "ModelRouter Desktop — embedded ModelRouter core with tray",
|
||||
"author": "ModelRouter",
|
||||
"main": "main.js",
|
||||
|
||||
@ -81,16 +81,42 @@ function adapter.transform_response(raw_body)
|
||||
|
||||
if resp.candidates and #resp.candidates > 0 then
|
||||
local cand = resp.candidates[1]
|
||||
local tools = {}
|
||||
if cand.content and cand.content.parts then
|
||||
for _, part in ipairs(cand.content.parts) do
|
||||
if part.text then
|
||||
unified.content = unified.content .. part.text
|
||||
elseif part.functionCall then
|
||||
-- Non-streaming tool calls used to be dropped here while
|
||||
-- transform_stream_chunk handled them, so a non-streaming
|
||||
-- agent turn looked like a plain text answer and the tool
|
||||
-- loop died. Gemini's args are already an object.
|
||||
local args = part.functionCall.args
|
||||
if type(args) == "string" then
|
||||
local aok, decoded = pcall(json.decode, args)
|
||||
args = aok and decoded or {}
|
||||
elseif type(args) ~= "table" then
|
||||
args = {}
|
||||
end
|
||||
table.insert(tools, {
|
||||
id = part.functionCall.id or ("call_" .. #tools),
|
||||
type = "function",
|
||||
name = part.functionCall.name or "",
|
||||
arguments = args
|
||||
})
|
||||
end
|
||||
end
|
||||
end
|
||||
if cand.finishReason then
|
||||
unified.finish_reason = cand.finishReason
|
||||
end
|
||||
if #tools > 0 then
|
||||
unified.tool_calls = tools
|
||||
-- Gemini reports finishReason "STOP" even when it emitted a
|
||||
-- functionCall; clients keyed on finish_reason would treat that as
|
||||
-- a completed answer and never run the tool.
|
||||
unified.finish_reason = "tool_calls"
|
||||
end
|
||||
end
|
||||
|
||||
return json.encode(unified)
|
||||
|
||||
@ -58,13 +58,39 @@ function adapter.transform_response(raw_body)
|
||||
local unified = {
|
||||
content = "",
|
||||
finish_reason = resp.done_reason or "",
|
||||
tool_calls = {},
|
||||
-- key must be token_usage to match Go's UnifiedResponse json tag
|
||||
token_usage = { prompt = p, completion = c, total = p + c }
|
||||
}
|
||||
|
||||
if resp.message then
|
||||
unified.content = resp.message.content or ""
|
||||
-- Non-streaming tool calls were previously dropped: the field was
|
||||
-- initialized to an empty table and never filled, while
|
||||
-- transform_stream_chunk handled them. A non-streaming agent turn thus
|
||||
-- looked like a plain answer and the tool loop stopped.
|
||||
if type(resp.message.tool_calls) == "table" and #resp.message.tool_calls > 0 then
|
||||
local tcs = {}
|
||||
for _, tc in ipairs(resp.message.tool_calls) do
|
||||
local fn = tc["function"] or {}
|
||||
-- Ollama sends arguments as an object already
|
||||
local args = fn.arguments
|
||||
if type(args) == "string" then
|
||||
local aok, decoded = pcall(json.decode, args)
|
||||
args = aok and decoded or {}
|
||||
elseif type(args) ~= "table" then
|
||||
args = {}
|
||||
end
|
||||
table.insert(tcs, {
|
||||
id = tc.id or ("call_" .. #tcs),
|
||||
type = tc.type or "function",
|
||||
name = fn.name or "",
|
||||
arguments = args
|
||||
})
|
||||
end
|
||||
unified.tool_calls = tcs
|
||||
-- Ollama reports done_reason "stop" alongside tool calls
|
||||
unified.finish_reason = "tool_calls"
|
||||
end
|
||||
end
|
||||
|
||||
return json.encode(unified)
|
||||
|
||||
@ -61,6 +61,34 @@ function adapter.transform_response(raw_body)
|
||||
-- sensenova-6.8-flash-lite 用 reasoning 字段而不是 reasoning_content
|
||||
unified.reasoning_content = ch.message.reasoning
|
||||
end
|
||||
-- Tool calls MUST be forwarded. Dropping them while keeping
|
||||
-- finish_reason="tool_calls" makes the client replay an assistant
|
||||
-- message whose function name/arguments are empty, and sensenova
|
||||
-- then rejects the next turn with
|
||||
-- 400 invalid tool_call function, function/name/arguments cannot be empty
|
||||
-- i.e. a tool-using conversation dies on its second request.
|
||||
if type(ch.message.tool_calls) == "table" and #ch.message.tool_calls > 0 then
|
||||
local tcs = {}
|
||||
for _, tc in ipairs(ch.message.tool_calls) do
|
||||
local fn = tc["function"] or {}
|
||||
-- arguments arrives as a JSON *string* on the wire; the
|
||||
-- unified shape expects a decoded object.
|
||||
local args = fn.arguments
|
||||
if type(args) == "string" then
|
||||
local args_ok, decoded = pcall(json.decode, args)
|
||||
args = args_ok and decoded or {}
|
||||
elseif type(args) ~= "table" then
|
||||
args = {}
|
||||
end
|
||||
table.insert(tcs, {
|
||||
id = tc.id,
|
||||
type = tc.type or "function",
|
||||
name = fn.name,
|
||||
arguments = args
|
||||
})
|
||||
end
|
||||
unified.tool_calls = tcs
|
||||
end
|
||||
end
|
||||
unified.finish_reason = ch.finish_reason or ""
|
||||
end
|
||||
|
||||
@ -25,6 +25,78 @@ function adapter.transform_request(raw_body)
|
||||
return json.encode(req)
|
||||
end
|
||||
|
||||
-- parse_text_tool_calls extracts tool calls that an upstream emitted as PLAIN
|
||||
-- TEXT instead of using the OpenAI tool_calls field.
|
||||
--
|
||||
-- trae-local-api's OpenAI endpoint (/v1/chat/completions) does not read the
|
||||
-- request's `tools` array at all, so the relayed model is never told the tool
|
||||
-- schema; it falls back to printing
|
||||
-- <tool_call>
|
||||
-- {"name": "get_weather", "arguments": {"location": "北京"}}
|
||||
-- </tool_call>
|
||||
-- into message.content, leaves message.tool_calls null, and reports
|
||||
-- finish_reason="stop". A client following the OpenAI contract therefore never
|
||||
-- sees a tool call: the agent loop terminates unexpectedly mid-conversation
|
||||
-- (and a hand-written replay produces an empty function name next turn).
|
||||
--
|
||||
-- Both tag spellings are accepted: the same codebase's Anthropic endpoint
|
||||
-- instructs models to emit <toolcall>, and models mix the two. Key names vary
|
||||
-- too (arguments / params / input), so all are tried.
|
||||
--
|
||||
-- Returns (tool_calls_array_or_nil, content_with_blocks_removed).
|
||||
local function parse_text_tool_calls(content)
|
||||
if type(content) ~= "string" or content == "" then return nil, content end
|
||||
if not (content:find("<tool_call", 1, true) or content:find("<toolcall", 1, true)) then
|
||||
return nil, content
|
||||
end
|
||||
|
||||
local tcs = {}
|
||||
local idx = 0
|
||||
|
||||
local function collect(pattern)
|
||||
for payload in content:gmatch(pattern) do
|
||||
local ok, obj = pcall(json.decode, payload)
|
||||
if ok and type(obj) == "table" then
|
||||
-- some models wrap it as {"function":{"name":..,"arguments":..}}
|
||||
local fn = obj["function"]
|
||||
local name = obj.name or (type(fn) == "table" and fn.name) or nil
|
||||
if name then
|
||||
local args = obj.arguments or obj.params or obj.input
|
||||
if args == nil and type(fn) == "table" then
|
||||
args = fn.arguments or fn.params
|
||||
end
|
||||
if type(args) == "string" then
|
||||
local aok, decoded = pcall(json.decode, args)
|
||||
args = aok and decoded or {}
|
||||
elseif type(args) ~= "table" then
|
||||
args = {}
|
||||
end
|
||||
idx = idx + 1
|
||||
table.insert(tcs, {
|
||||
id = obj.id or ("call_text_" .. idx),
|
||||
type = "function",
|
||||
name = name,
|
||||
arguments = args
|
||||
})
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
-- `<tag ...>` allows attributes; %s* handles the "</tool_call >" spacing
|
||||
-- that trae-local-api's own prompt example uses.
|
||||
collect("<tool_call[^>]*>%s*(.-)%s*</tool_call%s*>")
|
||||
collect("<toolcall[^>]*>%s*(.-)%s*</toolcall%s*>")
|
||||
|
||||
if #tcs == 0 then return nil, content end
|
||||
|
||||
-- drop the blocks from user-visible content; keep any surrounding prose
|
||||
local stripped = content:gsub("<tool_call[^>]*>%s*.-%s*</tool_call%s*>", "")
|
||||
stripped = stripped:gsub("<toolcall[^>]*>%s*.-%s*</toolcall%s*>", "")
|
||||
stripped = stripped:gsub("^%s+", ""):gsub("%s+$", "")
|
||||
return tcs, stripped
|
||||
end
|
||||
|
||||
function adapter.transform_response(raw_body)
|
||||
-- trae-local-api 偶发在非流式请求中返回 SSE 格式数据,
|
||||
-- 表现为多个 data: {...} 行或混合了 reasoning_chunk 等。
|
||||
@ -71,15 +143,21 @@ function adapter.transform_response(raw_body)
|
||||
if ch.message.reasoning_content then
|
||||
unified.reasoning_content = ch.message.reasoning_content
|
||||
end
|
||||
if type(ch.message.tool_calls) == "table" then
|
||||
if type(ch.message.tool_calls) == "table" and #ch.message.tool_calls > 0 then
|
||||
local tcs = {}
|
||||
for _, tc in ipairs(ch.message.tool_calls) do
|
||||
local args_ok, args = pcall(json.decode, tc["function"].arguments)
|
||||
if not args_ok then args = {} end
|
||||
local fn = tc["function"] or {}
|
||||
local args = fn.arguments
|
||||
if type(args) == "string" then
|
||||
local args_ok, decoded = pcall(json.decode, args)
|
||||
args = args_ok and decoded or {}
|
||||
elseif type(args) ~= "table" then
|
||||
args = {}
|
||||
end
|
||||
table.insert(tcs, {
|
||||
id = tc.id,
|
||||
type = tc.type or "function",
|
||||
name = tc["function"].name,
|
||||
name = fn.name,
|
||||
arguments = args
|
||||
})
|
||||
end
|
||||
@ -87,6 +165,22 @@ function adapter.transform_response(raw_body)
|
||||
end
|
||||
end
|
||||
unified.finish_reason = ch.finish_reason or ""
|
||||
|
||||
-- Fallback: trae-local-api's OpenAI endpoint drops the request's
|
||||
-- `tools` array entirely, so the relayed model is never told the tool
|
||||
-- schema and instead PRINTS a <tool_call>{...}</tool_call> block into
|
||||
-- content, leaving message.tool_calls null and finish_reason="stop".
|
||||
-- A client following the OpenAI contract then sees a normal completion
|
||||
-- and its agent loop terminates mid-conversation. Recover the
|
||||
-- structured call so the loop can continue.
|
||||
if unified.tool_calls == nil then
|
||||
local recovered, cleaned = parse_text_tool_calls(unified.content)
|
||||
if recovered then
|
||||
unified.tool_calls = recovered
|
||||
unified.content = cleaned
|
||||
unified.finish_reason = "tool_calls"
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
return json.encode(unified)
|
||||
|
||||
@ -1373,3 +1373,229 @@ func TestBurstLeftoverIsReclaimedUnderSteadyTraffic(t *testing.T) {
|
||||
created, idle)
|
||||
}
|
||||
}
|
||||
|
||||
// TestNonStreamToolCallsPreserved is the regression for the bug that killed
|
||||
// tool-using conversations on their SECOND request: several adapters handled
|
||||
// tool_calls in transform_stream_chunk but dropped them in
|
||||
// transform_response. A client then received finish_reason:"tool_calls" with no
|
||||
// tool_calls payload, replayed an assistant message whose function
|
||||
// name/arguments were empty, and the upstream rejected the next turn with
|
||||
//
|
||||
// 400 invalid tool_call function, function/name/arguments cannot be empty
|
||||
//
|
||||
// Every OpenAI-shaped adapter must forward non-streaming tool calls.
|
||||
func TestNonStreamToolCallsPreserved(t *testing.T) {
|
||||
vm := NewVM(freshAdapterDir(t))
|
||||
if err := vm.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer vm.Stop()
|
||||
|
||||
// upstream shape: arguments is a JSON *string*, as OpenAI sends it
|
||||
body := `{"choices":[{"index":0,"message":{"role":"assistant","content":"",` +
|
||||
`"tool_calls":[{"index":0,"id":"call_abc","type":"function",` +
|
||||
`"function":{"name":"get_weather","arguments":"{\"city\":\"Beijing\"}"}}]},` +
|
||||
`"finish_reason":"tool_calls"}],` +
|
||||
`"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15}}`
|
||||
|
||||
for _, name := range []string{
|
||||
"openai", "sensenova", "trae", "deepseek", "github",
|
||||
"groq", "kimicode", "mistral", "opencode", "agentrouter",
|
||||
} {
|
||||
out, err := vm.Transform(name, "transform_response", body)
|
||||
if err != nil {
|
||||
t.Fatalf("%s transform_response: %v", name, err)
|
||||
}
|
||||
var got struct {
|
||||
FinishReason string `json:"finish_reason"`
|
||||
ToolCalls []struct {
|
||||
ID string `json:"id"`
|
||||
Type string `json:"type"`
|
||||
Name string `json:"name"`
|
||||
Arguments map[string]interface{} `json:"arguments"`
|
||||
} `json:"tool_calls"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(out), &got); err != nil {
|
||||
t.Fatalf("%s: unmarshal %v (%s)", name, err, out)
|
||||
}
|
||||
if len(got.ToolCalls) != 1 {
|
||||
t.Fatalf("%s: dropped non-streaming tool_calls (client would replay an empty function): %s", name, out)
|
||||
}
|
||||
tc := got.ToolCalls[0]
|
||||
if tc.Name != "get_weather" {
|
||||
t.Errorf("%s: tool name lost: %s", name, out)
|
||||
}
|
||||
if tc.ID != "call_abc" {
|
||||
t.Errorf("%s: tool id lost: %s", name, out)
|
||||
}
|
||||
// arguments must be decoded into an object, not left as a string
|
||||
if tc.Arguments == nil || tc.Arguments["city"] != "Beijing" {
|
||||
t.Errorf("%s: arguments not decoded: %s", name, out)
|
||||
}
|
||||
if got.FinishReason != "tool_calls" {
|
||||
t.Errorf("%s: finish_reason lost: %s", name, out)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestGeminiNonStreamToolCalls covers Gemini's native shape (functionCall parts
|
||||
// under candidates[].content.parts). The streaming path handled these; the
|
||||
// non-streaming path silently returned text only.
|
||||
func TestGeminiNonStreamToolCalls(t *testing.T) {
|
||||
vm := NewVM(freshAdapterDir(t))
|
||||
if err := vm.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer vm.Stop()
|
||||
|
||||
body := `{"candidates":[{"content":{"parts":[` +
|
||||
`{"text":"checking"},` +
|
||||
`{"functionCall":{"name":"get_weather","args":{"city":"Beijing"}}}` +
|
||||
`]},"finishReason":"STOP"}],` +
|
||||
`"usageMetadata":{"promptTokenCount":8,"candidatesTokenCount":4,"totalTokenCount":12}}`
|
||||
|
||||
out, err := vm.Transform("gemini", "transform_response", body)
|
||||
if err != nil {
|
||||
t.Fatalf("gemini transform_response: %v", err)
|
||||
}
|
||||
var got struct {
|
||||
Content string `json:"content"`
|
||||
FinishReason string `json:"finish_reason"`
|
||||
ToolCalls []struct {
|
||||
Name string `json:"name"`
|
||||
Arguments map[string]interface{} `json:"arguments"`
|
||||
} `json:"tool_calls"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(out), &got); err != nil {
|
||||
t.Fatalf("unmarshal: %v (%s)", err, out)
|
||||
}
|
||||
if len(got.ToolCalls) != 1 || got.ToolCalls[0].Name != "get_weather" {
|
||||
t.Fatalf("gemini dropped non-streaming functionCall: %s", out)
|
||||
}
|
||||
if got.ToolCalls[0].Arguments["city"] != "Beijing" {
|
||||
t.Errorf("gemini args lost: %s", out)
|
||||
}
|
||||
// a tool call must not be reported as a plain stop
|
||||
if got.FinishReason != "tool_calls" {
|
||||
t.Errorf("gemini finish_reason should become tool_calls: %s", out)
|
||||
}
|
||||
}
|
||||
|
||||
// TestOllamaNonStreamToolCalls covers Ollama's non-streaming shape
|
||||
// (message.tool_calls with an already-decoded arguments object). The adapter
|
||||
// initialized tool_calls to an empty table and never filled it.
|
||||
func TestOllamaNonStreamToolCalls(t *testing.T) {
|
||||
vm := NewVM(freshAdapterDir(t))
|
||||
if err := vm.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer vm.Stop()
|
||||
|
||||
body := `{"message":{"role":"assistant","content":"",` +
|
||||
`"tool_calls":[{"function":{"name":"get_weather","arguments":{"city":"Beijing"}}}]},` +
|
||||
`"done_reason":"stop","prompt_eval_count":7,"eval_count":3}`
|
||||
|
||||
out, err := vm.Transform("ollama", "transform_response", body)
|
||||
if err != nil {
|
||||
t.Fatalf("ollama transform_response: %v", err)
|
||||
}
|
||||
var got struct {
|
||||
FinishReason string `json:"finish_reason"`
|
||||
ToolCalls []struct {
|
||||
Name string `json:"name"`
|
||||
Arguments map[string]interface{} `json:"arguments"`
|
||||
} `json:"tool_calls"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(out), &got); err != nil {
|
||||
t.Fatalf("unmarshal: %v (%s)", err, out)
|
||||
}
|
||||
if len(got.ToolCalls) != 1 || got.ToolCalls[0].Name != "get_weather" {
|
||||
t.Fatalf("ollama dropped non-streaming tool_calls: %s", out)
|
||||
}
|
||||
if got.ToolCalls[0].Arguments["city"] != "Beijing" {
|
||||
t.Errorf("ollama args lost: %s", out)
|
||||
}
|
||||
if got.FinishReason != "tool_calls" {
|
||||
t.Errorf("ollama finish_reason should become tool_calls: %s", out)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTraeTextToolCallRecovery covers the trae-specific failure: its upstream
|
||||
// (trae-local-api's OpenAI endpoint) never reads the request's `tools` array, so
|
||||
// the relayed model prints a <tool_call> block as TEXT, leaves
|
||||
// message.tool_calls null and reports finish_reason:"stop". An OpenAI client
|
||||
// then sees an ordinary completion and its agent loop ends mid-conversation.
|
||||
// The adapter must recover the structured call, strip the block from content,
|
||||
// and correct finish_reason.
|
||||
func TestTraeTextToolCallRecovery(t *testing.T) {
|
||||
vm := NewVM(freshAdapterDir(t))
|
||||
if err := vm.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer vm.Stop()
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
body string
|
||||
}{
|
||||
{
|
||||
"tool_call tag with arguments",
|
||||
`{"choices":[{"message":{"role":"assistant","content":"<tool_call>\n{\"name\": \"get_weather\", \"arguments\": {\"city\": \"Beijing\"}}\n</tool_call>"},"finish_reason":"stop"}]}`,
|
||||
},
|
||||
{
|
||||
// the same codebase's Anthropic endpoint prompts for <toolcall> and "params"
|
||||
"toolcall tag with params",
|
||||
`{"choices":[{"message":{"role":"assistant","content":"<toolcall>{\"name\": \"get_weather\", \"params\": {\"city\": \"Beijing\"}}</toolcall>"},"finish_reason":"stop"}]}`,
|
||||
},
|
||||
{
|
||||
"prose around the block is preserved",
|
||||
`{"choices":[{"message":{"role":"assistant","content":"Let me check.\n<tool_call>{\"name\":\"get_weather\",\"arguments\":{\"city\":\"Beijing\"}}</tool_call>"},"finish_reason":"stop"}]}`,
|
||||
},
|
||||
}
|
||||
|
||||
for _, c := range cases {
|
||||
out, err := vm.Transform("trae", "transform_response", c.body)
|
||||
if err != nil {
|
||||
t.Fatalf("%s: transform: %v", c.name, err)
|
||||
}
|
||||
var got struct {
|
||||
Content string `json:"content"`
|
||||
FinishReason string `json:"finish_reason"`
|
||||
ToolCalls []struct {
|
||||
Name string `json:"name"`
|
||||
Arguments map[string]interface{} `json:"arguments"`
|
||||
} `json:"tool_calls"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(out), &got); err != nil {
|
||||
t.Fatalf("%s: unmarshal %v (%s)", c.name, err, out)
|
||||
}
|
||||
if len(got.ToolCalls) != 1 {
|
||||
t.Fatalf("%s: text tool call not recovered: %s", c.name, out)
|
||||
}
|
||||
if got.ToolCalls[0].Name != "get_weather" {
|
||||
t.Errorf("%s: name wrong: %s", c.name, out)
|
||||
}
|
||||
if got.ToolCalls[0].Arguments["city"] != "Beijing" {
|
||||
t.Errorf("%s: args wrong: %s", c.name, out)
|
||||
}
|
||||
if got.FinishReason != "tool_calls" {
|
||||
t.Errorf("%s: finish_reason must be corrected to tool_calls: %s", c.name, out)
|
||||
}
|
||||
if strings.Contains(got.Content, "<tool_call") || strings.Contains(got.Content, "<toolcall") {
|
||||
t.Errorf("%s: raw block left in user-visible content: %s", c.name, out)
|
||||
}
|
||||
}
|
||||
|
||||
// a plain text answer must be untouched
|
||||
plain := `{"choices":[{"message":{"role":"assistant","content":"just text"},"finish_reason":"stop"}]}`
|
||||
out, err := vm.Transform("trae", "transform_response", plain)
|
||||
if err != nil {
|
||||
t.Fatalf("plain: %v", err)
|
||||
}
|
||||
if strings.Contains(out, "tool_calls") {
|
||||
t.Errorf("plain answer must not gain tool_calls: %s", out)
|
||||
}
|
||||
if !strings.Contains(out, `"finish_reason":"stop"`) {
|
||||
t.Errorf("plain answer finish_reason changed: %s", out)
|
||||
}
|
||||
}
|
||||
|
||||
@ -34,9 +34,14 @@ import (
|
||||
// cooldown always expires and any success resets the state, so a fixed
|
||||
// upstream recovers on its own.
|
||||
type ModelState struct {
|
||||
pref atomic.Int64 // +1 per success / -5 per failure, clamped
|
||||
pref atomic.Int64 // +prefSuccessStep per success / -prefFailStep per failure, clamped
|
||||
failCount atomic.Int64
|
||||
cooldownUntil atomic.Int64 // unix seconds; 0 = schedulable
|
||||
// prefTouched is the unix second pref was last moved by an outcome. It is
|
||||
// the reference point for decay: a negative score with no traffic since
|
||||
// prefDecayAfter drifts one step toward 0 per elapsed interval, so a slot
|
||||
// cannot stay written off forever just because nothing selected it.
|
||||
prefTouched atomic.Int64
|
||||
// cooldownFrom is the unix second the current cooldown window opened. It
|
||||
// is the reference point for the half-cooldown probe gate: probing is
|
||||
// allowed only in the second half of [cooldownFrom, cooldownUntil), so a
|
||||
@ -50,10 +55,35 @@ type ModelState struct {
|
||||
}
|
||||
|
||||
const (
|
||||
// prefFailStep is the penalty for a REAL failure (5xx, transport error,
|
||||
// unusable adapter output): something is wrong with the upstream or the
|
||||
// model, so two of them are enough to push a slot to the back of its tier.
|
||||
prefFailStep = 5
|
||||
prefMin = -20
|
||||
prefMax = 20
|
||||
backoffBase = 5 * time.Second
|
||||
// prefQuotaStep is the penalty for running out of allowance (quota/TPM
|
||||
// exhausted). Deliberately much smaller than prefFailStep: an exhausted
|
||||
// quota means "no budget right now", not "this model is bad", and for
|
||||
// allowance-metered sources it is an everyday event rather than a fault.
|
||||
// Charging it the full failure penalty made such sources drift to a deeply
|
||||
// negative score with ZERO visible failures - RecordQuotaExhausted
|
||||
// intentionally does not touch failCount, so the WebUI showed fails=0 and
|
||||
// cooling=false while the score sat at -10.
|
||||
prefQuotaStep = 1
|
||||
// prefSuccessStep is the reward for a successful request. Recovery used to
|
||||
// be +1 against a -5 penalty, a 5:1 asymmetry that made a slot at -10 need
|
||||
// ten consecutive successes just to reach neutral. Rewarding +2 keeps
|
||||
// failures meaningful (a real failure still costs more than one success
|
||||
// earns) while letting a recovered slot return to normal rotation quickly.
|
||||
prefSuccessStep = 2
|
||||
prefMin = -20
|
||||
prefMax = 20
|
||||
// prefDecayAfter is how long a slot may sit at a negative score without any
|
||||
// traffic before the score drifts back toward neutral. Without this a slot
|
||||
// that was penalised and then simply not selected again (the score itself
|
||||
// makes it unattractive) stays penalised indefinitely: it has no way to earn
|
||||
// the successes that would rehabilitate it. Applied lazily on read, so there
|
||||
// are no timers.
|
||||
prefDecayAfter = 2 * time.Minute
|
||||
backoffBase = 5 * time.Second
|
||||
// backoffCap bounds the exponential schedule. It used to be 30min, which
|
||||
// in practice meant "a slot whose upstream blipped is unusable for half an
|
||||
// hour even though its quota recovered in seconds". Combined with the
|
||||
@ -142,8 +172,55 @@ func (s *ModelState) RecordFailure(auth bool) {
|
||||
if auth {
|
||||
pen *= 2
|
||||
}
|
||||
s.pref.Add(-pen)
|
||||
s.addPref(-pen)
|
||||
}
|
||||
|
||||
// addPref moves the preference score and stamps prefTouched, which restarts the
|
||||
// decay clock. Every outcome-driven score change must go through here so decay
|
||||
// never fires on a slot that is actively being scored.
|
||||
func (s *ModelState) addPref(delta int64) {
|
||||
s.pref.Add(delta)
|
||||
clampPref(&s.pref, prefMin, prefMax)
|
||||
s.prefTouched.Store(time.Now().Unix())
|
||||
}
|
||||
|
||||
// decayPref drifts a negative score back toward neutral once a slot has been
|
||||
// idle for prefDecayAfter. This closes a starvation loop: a penalised slot is
|
||||
// unattractive, so it stops being selected, so it never earns the successes that
|
||||
// would rehabilitate it, so it stays penalised. One step per elapsed interval,
|
||||
// never past 0, and positive scores are left alone (a proven-good slot should
|
||||
// not be dragged down for being idle).
|
||||
//
|
||||
// Called lazily from Pref(); no timers and no goroutines.
|
||||
func (s *ModelState) decayPref() {
|
||||
cur := s.pref.Load()
|
||||
if cur >= 0 {
|
||||
return
|
||||
}
|
||||
last := s.prefTouched.Load()
|
||||
if last <= 0 {
|
||||
// no reference point yet: start the clock instead of decaying blindly
|
||||
s.prefTouched.CompareAndSwap(0, time.Now().Unix())
|
||||
return
|
||||
}
|
||||
step := int64(prefDecayAfter / time.Second)
|
||||
if step <= 0 {
|
||||
return
|
||||
}
|
||||
elapsed := time.Now().Unix() - last
|
||||
if elapsed < step {
|
||||
return
|
||||
}
|
||||
steps := elapsed / step
|
||||
target := cur + steps
|
||||
if target > 0 {
|
||||
target = 0
|
||||
}
|
||||
// Advance prefTouched by exactly the consumed intervals so leftover time
|
||||
// still counts toward the next step.
|
||||
if s.pref.CompareAndSwap(cur, target) {
|
||||
s.prefTouched.Store(last + steps*step)
|
||||
}
|
||||
}
|
||||
|
||||
// openCooldown starts a fresh cooldown window of length cd. Recording both
|
||||
@ -164,8 +241,10 @@ func (s *ModelState) openCooldown(cd time.Duration) {
|
||||
func (s *ModelState) RecordRateLimit() {
|
||||
s.failCount.Add(1)
|
||||
s.openCooldown(rateLimitCooldown)
|
||||
s.pref.Add(-prefFailStep)
|
||||
clampPref(&s.pref, prefMin, prefMax)
|
||||
// A 429 is "too fast", not "broken": charge it like an allowance event, not
|
||||
// like a fault, or a throttled-but-healthy source sinks purely for being
|
||||
// popular.
|
||||
s.addPref(-prefQuotaStep)
|
||||
}
|
||||
|
||||
// RecordQuotaExhausted marks the model as out of allowance until its quota
|
||||
@ -185,20 +264,24 @@ func (s *ModelState) RecordQuotaExhausted(resetIn time.Duration) {
|
||||
// No failCount bump: an exhausted quota is not an error streak, and
|
||||
// letting it inflate failCount would make the NEXT real failure jump
|
||||
// straight to the capped backoff.
|
||||
s.pref.Add(-prefFailStep)
|
||||
clampPref(&s.pref, prefMin, prefMax)
|
||||
//
|
||||
// The score penalty is prefQuotaStep, not prefFailStep, for the same
|
||||
// reason: cooldown alone already keeps the slot out of rotation until the
|
||||
// window resets, so the score only needs to express a mild preference for
|
||||
// slots with budget left. Charging the full failure penalty here is what
|
||||
// drove allowance-metered sources to -10 and beyond with fails=0.
|
||||
s.addPref(-prefQuotaStep)
|
||||
}
|
||||
|
||||
// RecordSuccess resets the failure counter and cooldown and bumps the
|
||||
// preference score by one. A single success is enough to lift a slot off the
|
||||
// prefMin floor (prefMin+1 > prefMin), which is what lets one successful probe
|
||||
// return a written-off slot to normal rotation.
|
||||
// RecordSuccess resets the failure counter and cooldown and rewards the
|
||||
// preference score. A single success lifts a slot off the prefMin floor, which
|
||||
// is what lets one successful probe return a written-off slot to normal
|
||||
// rotation.
|
||||
func (s *ModelState) RecordSuccess() {
|
||||
s.failCount.Store(0)
|
||||
s.cooldownUntil.Store(0)
|
||||
s.cooldownFrom.Store(0)
|
||||
s.pref.Add(1)
|
||||
clampPref(&s.pref, prefMin, prefMax)
|
||||
s.addPref(prefSuccessStep)
|
||||
}
|
||||
|
||||
func (s *ModelState) reset() {
|
||||
@ -207,6 +290,7 @@ func (s *ModelState) reset() {
|
||||
s.cooldownFrom.Store(0)
|
||||
s.probeInFlight.Store(false)
|
||||
s.pref.Store(0)
|
||||
s.prefTouched.Store(0)
|
||||
}
|
||||
|
||||
// TryProbe attempts to claim the single probe permit for a model that normal
|
||||
@ -225,6 +309,7 @@ func (s *ModelState) TryProbe() bool {
|
||||
now := time.Now().Unix()
|
||||
until := s.cooldownUntil.Load()
|
||||
if until <= now {
|
||||
s.decayPref() // a naturally-recovered slot needs no probe
|
||||
if s.pref.Load() > prefMin {
|
||||
return false // schedulable normally, no probe needed
|
||||
}
|
||||
@ -251,8 +336,13 @@ func (s *ModelState) Probing() bool { return s.probeInFlight.Load() }
|
||||
// when the model is not cooling.
|
||||
func (s *ModelState) CooldownFrom() int64 { return s.cooldownFrom.Load() }
|
||||
|
||||
// Pref is the current preference score (higher = preferred).
|
||||
func (s *ModelState) Pref() int64 { return s.pref.Load() }
|
||||
// Pref is the current preference score (higher = preferred). Reading it also
|
||||
// applies idle decay, so a slot penalised long ago and never selected since
|
||||
// drifts back toward neutral instead of being written off permanently.
|
||||
func (s *ModelState) Pref() int64 {
|
||||
s.decayPref()
|
||||
return s.pref.Load()
|
||||
}
|
||||
|
||||
// FailCount is the number of consecutive failures.
|
||||
func (s *ModelState) FailCount() int64 { return s.failCount.Load() }
|
||||
|
||||
@ -242,8 +242,10 @@ func TestModelStateCooldownAndRecovery(t *testing.T) {
|
||||
if st.FailCount() != 0 {
|
||||
t.Fatalf("fail count after success = %d", st.FailCount())
|
||||
}
|
||||
if st.Pref() != 1-int64(prefFailStep) {
|
||||
t.Fatalf("pref after one failure (-5) then success (+1) = %d, want %d", st.Pref(), 1-int64(prefFailStep))
|
||||
// failure penalties prefFailStep, success rewards prefSuccessStep:
|
||||
// -5 then +2 = -3
|
||||
if st.Pref() != int64(prefSuccessStep)-int64(prefFailStep) {
|
||||
t.Fatalf("pref after one failure then success = %d, want %d", st.Pref(), int64(prefSuccessStep)-int64(prefFailStep))
|
||||
}
|
||||
}
|
||||
|
||||
@ -265,8 +267,8 @@ func TestModelStateRateLimitShortCooldown(t *testing.T) {
|
||||
if st.FailCount() != 15 {
|
||||
t.Fatalf("fail count = %d, want 15 (counted but not escalating)", st.FailCount())
|
||||
}
|
||||
if st.Pref() != -15*int64(prefFailStep) && st.Pref() > int64(prefMin) {
|
||||
t.Fatalf("pref = %d", st.Pref())
|
||||
if st.Pref() != -15*int64(prefQuotaStep) {
|
||||
t.Fatalf("pref after 15 x 429 = %d, want %d (a 429 is a quota-class event, not a fault)", st.Pref(), -15*int64(prefQuotaStep))
|
||||
}
|
||||
if p.ModelAvailable("m") {
|
||||
t.Fatal("model must be cooling right after a 429")
|
||||
@ -1044,3 +1046,140 @@ func TestRecoveryProbeBudgetIsBounded(t *testing.T) {
|
||||
t.Fatalf("3 probe windows must cost exactly 3 upstream requests, got %d", extra)
|
||||
}
|
||||
}
|
||||
|
||||
// TestQuotaPenaltyIsLighterThanFailure is the regression for the report
|
||||
// "今天一次失败没看到,分数目前已经 -10 了" on an allowance-metered source
|
||||
// (sensenova). RecordQuotaExhausted deliberately does not touch failCount — the
|
||||
// WebUI therefore showed fails=0 and cooling=false — yet it used to charge the
|
||||
// SAME score penalty as a real failure, so two ordinary quota resets dragged the
|
||||
// slot to -10 with no visible fault. An exhausted allowance means "no budget
|
||||
// right now", not "this model is broken": the cooldown already keeps it out of
|
||||
// rotation, so the score penalty must be mild.
|
||||
func TestQuotaPenaltyIsLighterThanFailure(t *testing.T) {
|
||||
quota := &ModelState{}
|
||||
quota.RecordQuotaExhausted(time.Minute)
|
||||
fail := &ModelState{}
|
||||
fail.RecordFailure(false)
|
||||
|
||||
qp, fp := quota.Pref(), fail.Pref()
|
||||
if qp <= fp {
|
||||
t.Fatalf("a quota reset must cost less score than a real failure: quota=%d failure=%d", qp, fp)
|
||||
}
|
||||
if qp != -prefQuotaStep {
|
||||
t.Errorf("quota penalty = %d, want -%d", qp, prefQuotaStep)
|
||||
}
|
||||
if fp != -prefFailStep {
|
||||
t.Errorf("failure penalty = %d, want -%d", fp, prefFailStep)
|
||||
}
|
||||
|
||||
// The reported symptom: several quota resets in a day must NOT approach the
|
||||
// prefMin floor, since each one is a normal event for a metered source.
|
||||
s := &ModelState{}
|
||||
for i := 0; i < 5; i++ {
|
||||
s.RecordQuotaExhausted(time.Minute)
|
||||
}
|
||||
if got := s.Pref(); got <= prefMin/2 {
|
||||
t.Errorf("5 quota resets sank the score to %d (floor %d); a metered source would be written off", got, prefMin)
|
||||
}
|
||||
if s.FailCount() != 0 {
|
||||
t.Errorf("quota exhaustion must not inflate failCount, got %d", s.FailCount())
|
||||
}
|
||||
}
|
||||
|
||||
// TestRateLimitPenaltyIsMild covers the sibling case: a 429 means "too fast",
|
||||
// not "broken", so a popular-but-healthy slot must not sink for being throttled.
|
||||
func TestRateLimitPenaltyIsMild(t *testing.T) {
|
||||
s := &ModelState{}
|
||||
s.RecordRateLimit()
|
||||
if got := s.Pref(); got != -prefQuotaStep {
|
||||
t.Errorf("429 penalty = %d, want -%d (same class as a quota event)", got, prefQuotaStep)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSuccessRecoversFasterThanBefore pins the penalty/reward asymmetry. Reward
|
||||
// used to be +1 against a -5 penalty, so a slot at -10 needed TEN consecutive
|
||||
// successes just to reach neutral — in practice it never got there, because a
|
||||
// low score makes the slot unattractive to the scheduler in the first place.
|
||||
func TestSuccessRecoversFasterThanBefore(t *testing.T) {
|
||||
s := &ModelState{}
|
||||
s.RecordFailure(false)
|
||||
s.RecordFailure(false) // -10, the score from the bug report
|
||||
|
||||
n := 0
|
||||
for s.Pref() < 0 {
|
||||
s.RecordSuccess()
|
||||
n++
|
||||
if n > 50 {
|
||||
t.Fatal("score never recovered")
|
||||
}
|
||||
}
|
||||
// with +2 per success, -10 needs 5; the old +1 needed 10
|
||||
if n > 5 {
|
||||
t.Errorf("recovery from -10 took %d successes, want <= 5", n)
|
||||
}
|
||||
// a real failure must still outweigh a single success, or failures stop
|
||||
// meaning anything
|
||||
if prefSuccessStep >= prefFailStep {
|
||||
t.Errorf("success reward %d must stay below the failure penalty %d", prefSuccessStep, prefFailStep)
|
||||
}
|
||||
}
|
||||
|
||||
// TestPrefDecayRehabilitatesIdleSlot closes the starvation loop: a penalised
|
||||
// slot is unattractive, so the scheduler stops picking it, so it never earns the
|
||||
// successes that would rehabilitate it. Without decay the score stays negative
|
||||
// forever even though the upstream may be perfectly healthy.
|
||||
func TestPrefDecayRehabilitatesIdleSlot(t *testing.T) {
|
||||
s := &ModelState{}
|
||||
s.RecordFailure(false)
|
||||
s.RecordFailure(false)
|
||||
if s.Pref() != -2*prefFailStep {
|
||||
t.Fatalf("setup: pref=%d", s.Pref())
|
||||
}
|
||||
|
||||
// backdate the last touch by three decay intervals
|
||||
step := int64(prefDecayAfter / time.Second)
|
||||
s.prefTouched.Store(time.Now().Unix() - 3*step)
|
||||
|
||||
got := s.Pref()
|
||||
if got != -2*prefFailStep+3 {
|
||||
t.Errorf("after 3 idle intervals pref=%d, want %d", got, -2*prefFailStep+3)
|
||||
}
|
||||
|
||||
// decay must never overshoot into positive territory
|
||||
s.prefTouched.Store(time.Now().Unix() - 1000*step)
|
||||
if got := s.Pref(); got != 0 {
|
||||
t.Errorf("decay overshot to %d, want to stop at 0", got)
|
||||
}
|
||||
|
||||
// a positive score must not be dragged down for being idle
|
||||
good := &ModelState{}
|
||||
good.RecordSuccess()
|
||||
good.RecordSuccess()
|
||||
before := good.Pref()
|
||||
good.prefTouched.Store(time.Now().Unix() - 1000*step)
|
||||
if after := good.Pref(); after != before {
|
||||
t.Errorf("idle decay must not touch a positive score: %d -> %d", before, after)
|
||||
}
|
||||
}
|
||||
|
||||
// TestDecayLiftsSlotOffTheFloor verifies decay actually restores schedulability:
|
||||
// a slot at prefMin is refused by ModelAvailable, and before decay existed only
|
||||
// a probe could rescue it.
|
||||
func TestDecayLiftsSlotOffTheFloor(t *testing.T) {
|
||||
s := &ModelState{}
|
||||
for i := 0; i < 10; i++ {
|
||||
s.RecordFailure(false)
|
||||
}
|
||||
if s.Pref() != prefMin {
|
||||
t.Fatalf("setup: want floor %d, got %d", prefMin, s.Pref())
|
||||
}
|
||||
// clear the cooldown so only the score gates the slot
|
||||
s.cooldownUntil.Store(0)
|
||||
s.cooldownFrom.Store(0)
|
||||
|
||||
step := int64(prefDecayAfter / time.Second)
|
||||
s.prefTouched.Store(time.Now().Unix() - 2*step)
|
||||
if got := s.Pref(); got <= prefMin {
|
||||
t.Errorf("decay failed to lift the slot off the floor: %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user