diff --git a/internal/backup/backup.go b/internal/backup/backup.go index 3b9ecbd4..4b7e797f 100644 --- a/internal/backup/backup.go +++ b/internal/backup/backup.go @@ -230,7 +230,7 @@ var secretHeader = regexp.MustCompile(`(?i)auth|key|token|secret|cookie|session| func Secret(name string) bool { return secretHeader.MatchString(name) } func withoutKeys(p provider.Provider) provider.Provider { - p.Key, p.KeyName, p.Keys, p.KeyProtocol = "", "", nil, "" + p.Key, p.KeyName, p.Keys, p.KeyProtocol, p.KeyWeight = "", "", nil, "", 0 p.BalanceToken = "" if len(p.Headers) > 0 { h := map[string]string{} diff --git a/internal/davsync/davsync.go b/internal/davsync/davsync.go index a162af2e..48953362 100644 --- a/internal/davsync/davsync.go +++ b/internal/davsync/davsync.go @@ -644,7 +644,7 @@ func take(to *backup.Bundle, from backup.Bundle, part string) { ps = slices.Clone(ps) for i, p := range ps { if k, ok := keys[p.ID]; ok && p.Key == "" && len(p.Keys) == 0 { - ps[i].Key, ps[i].KeyName, ps[i].Keys, ps[i].KeyProtocol = k.Key, k.KeyName, k.Keys, k.KeyProtocol + ps[i].Key, ps[i].KeyName, ps[i].Keys, ps[i].KeyProtocol, ps[i].KeyWeight = k.Key, k.KeyName, k.Keys, k.KeyProtocol, k.KeyWeight if p.BalanceToken == "" { ps[i].BalanceToken = k.BalanceToken } diff --git a/internal/gateway/affinity.go b/internal/gateway/affinity.go index 026d797b..37e9dc0d 100644 --- a/internal/gateway/affinity.go +++ b/internal/gateway/affinity.go @@ -281,6 +281,7 @@ func affine(scope, mode string, rotate bool, in http.Header, from provider.Proto case "session", "turn", "cache": a.Kept = true if at > 0 { + weightedKept(cs[0], cs[at]) cs = append(append([]candidate{cs[at]}, cs[:at]...), cs[at+1:]...) order := append(append([]Weighed{pl.order[at]}, pl.order[:at]...), pl.order[at+1:]...) pl.order = order diff --git a/internal/gateway/fallback.go b/internal/gateway/fallback.go index 2c274e9a..55f2c97b 100644 --- a/internal/gateway/fallback.go +++ b/internal/gateway/fallback.go @@ -302,7 +302,7 @@ func (s *Server) plan(p provider.Provider, model string, from provider.Protocol) cs, wg := weigh(q, cs, m, from) for i, c := range cs { w := weighed(c, q, wg, fallback, from) - w.Turn = i == 0 && q.Routing == provider.Rotate && len(cs) > 1 + w.Turn = i == 0 && (q.Routing == provider.Rotate || q.Routing == provider.Weighted) && len(cs) > 1 pl.order = append(pl.order, w) } pl.order = append(pl.order, asideOf(aside, q, fallback, from)...) @@ -464,7 +464,7 @@ func planLevel(g provider.Group, ms []provider.Member, depth int, from provider. cs, wg := weigh(m.Provider, keys(m), m.Model, from) for k, c := range cs { w := weighed(c, m.Provider, wg, false, from) - w.Turn = k == 0 && m.Provider.Routing == provider.Rotate && len(cs) > 1 + w.Turn = k == 0 && (m.Provider.Routing == provider.Rotate || m.Provider.Routing == provider.Weighted) && len(cs) > 1 w.Via = m.Groups() pl.order = append(pl.order, w) } diff --git a/internal/gateway/routing.go b/internal/gateway/routing.go index 0007d063..3dca772e 100644 --- a/internal/gateway/routing.go +++ b/internal/gateway/routing.go @@ -705,6 +705,8 @@ func weighRouted(p provider.Provider, cs []candidate, model string, from provide routed.turn[p.ID]++ routed.Unlock() cs = append(append([]candidate{}, cs[n:]...), cs[:n]...) + case provider.Weighted: + cs = weightedFirst(p.ID, cs) case provider.LeastUsed: // a subscription by the share of its allowance used, as the vendor // says; then, and for keys, by what magpie sent it lately diff --git a/internal/gateway/sink.go b/internal/gateway/sink.go index bb971cea..7bbc5c2a 100644 --- a/internal/gateway/sink.go +++ b/internal/gateway/sink.go @@ -54,9 +54,10 @@ func sunkAt(c candidate) (sinking, bool) { } // sinks says whether an order routed so sends the rate limited back: -// in turn goes round already, and a manual group has one member. +// in turn and by weight go round already, and a manual group has one +// member. func sinks(on bool, routing string) bool { - return on && routing != provider.Rotate && routing != provider.Manual + return on && routing != provider.Rotate && routing != provider.Weighted && routing != provider.Manual } // sinkOrder is the order of n candidates with those that sank moved diff --git a/internal/gateway/weighted.go b/internal/gateway/weighted.go new file mode 100644 index 00000000..f7ff4809 --- /dev/null +++ b/internal/gateway/weighted.go @@ -0,0 +1,108 @@ +package gateway + +// Weighted routing (#841): a provider's requests spread over its keys by +// each key's weight, as NGINX's smooth weighted round robin does — keys +// weighing 3 and 1 go a, a, b, a, a, a, b, a…, the ratio held over a +// handful of requests rather than only on average. A key resting after a +// failure is left out of the round while it rests, so the others share +// its requests by their own weights and it takes up its share when it is +// back. A conversation still stays with the key that answered it while +// its affinity says (affinity.go); a request kept so counts against the +// key that took it, not the one the round picked (weightedKept), so new +// conversations make up the share. + +import ( + "sync" + + "github.com/yetone/magpie/internal/provider" +) + +// wrr is each weighted provider's round: every key's current weight, and +// the total the last pick was charged. +var wrr = struct { + sync.Mutex + m map[string]*wrrRound // provider id → its round +}{m: map[string]*wrrRound{}} + +type wrrRound struct { + cur map[string]int // candidate rest → current weight + total int +} + +// keyWeight is a candidate's weight: its key's, 1 when none. +func keyWeight(c candidate) int { return max(c.p.KeyWeight, 1) } + +// resting says whether c rests now, as restLast will find. +func resting(c candidate) bool { + if _, ok := restOf(c.restKey()); ok { + return true + } + if c.restID() != c.restKey() { + _, ok := restOf(c.restID()) + return ok + } + return false +} + +// weightedFirst puts first the key the provider id's round picks now, the +// others after it in their order, for failover. +func weightedFirst(id string, cs []candidate) []candidate { + wrr.Lock() + defer wrr.Unlock() + r := wrr.m[id] + if r == nil { + r = &wrrRound{cur: map[string]int{}} + wrr.m[id] = r + } + // keys no longer on are dropped from the round + on := make(map[string]bool, len(cs)) + for _, c := range cs { + on[c.rest] = true + } + for k := range r.cur { + if !on[k] { + delete(r.cur, k) + } + } + best, total := -1, 0 + for i, c := range cs { + if resting(c) { + continue + } + w := keyWeight(c) + r.cur[c.rest] += w + total += w + if best < 0 || r.cur[c.rest] > r.cur[cs[best].rest] { + best = i + } + } + if best < 0 { + return cs // all resting: restLast orders them + } + r.cur[cs[best].rest] -= total + r.total = total + if best == 0 { + return cs + } + return append(append([]candidate{cs[best]}, cs[:best]...), cs[best+1:]...) +} + +// weightedKept moves the pick of picked's round to kept, when affinity +// keeps a conversation with kept, another key of the same weighted +// provider: the round counts the requests each key really took. +func weightedKept(picked, kept candidate) { + if picked.p.Routing != provider.Weighted || picked.p.ID != kept.p.ID || picked.rest == kept.rest { + return + } + wrr.Lock() + defer wrr.Unlock() + r := wrr.m[picked.p.ID] + if r == nil { + return + } + if _, ok := r.cur[kept.rest]; !ok { + return + } + r.cur[picked.rest] += r.total + r.cur[kept.rest] -= r.total +} diff --git a/internal/gateway/weighted_keys_test.go b/internal/gateway/weighted_keys_test.go new file mode 100644 index 00000000..a6637057 --- /dev/null +++ b/internal/gateway/weighted_keys_test.go @@ -0,0 +1,148 @@ +package gateway + +import ( + "strings" + "testing" + "time" + + "github.com/yetone/magpie/internal/provider" +) + +// Weighted routing (#841): a provider's keys weighing 3 and 1 take three +// requests to one, held over every few requests rather than only on +// average, and a key resting is left out of the round, the others taking +// its share by their weights. +func TestWeightedKeys(t *testing.T) { + p := provider.Provider{ID: "wk", Chat: "https://wk.example/v1", Key: "ka", KeyName: "team-a", KeyWeight: 3, + Keys: []provider.KeyAccount{{Name: "team-b", Key: "kb"}, {Name: "team-c", Key: "kc", Weight: 4}, {Name: "spare", Key: "kd", Off: true, Weight: 9}}, + Routing: provider.Weighted} + a, b, c := "wk#"+provider.KeyID("ka"), "wk#"+provider.KeyID("kb"), "wk#"+provider.KeyID("kc") + t.Cleanup(func() { + wrr.Lock() + delete(wrr.m, p.ID) + wrr.Unlock() + restingUntil.Lock() + delete(restingUntil.m, c) + restingUntil.Unlock() + }) + pick := func() string { + cs := route(p, perKey(p, "m", provider.Chat), "m", provider.Chat) + if len(cs) != 3 { + t.Fatalf("%d candidates: %s", len(cs), restsOf(cs)) + } + return cs[0].rest + } + // 3 : 1 : 4 over every round of 8, smoothly: never the same key more + // than twice running when another weighs as much + got := map[string]int{} + var seq []string + for range 8 * 50 { + k := pick() + got[k]++ + seq = append(seq, k) + if len(seq)%8 == 0 { + round := map[string]int{} + for _, s := range seq[len(seq)-8:] { + round[s]++ + } + if round[a] != 3 || round[b] != 1 || round[c] != 4 { + t.Fatalf("round %d: a %d b %d c %d, want 3 1 4", len(seq)/8, round[a], round[b], round[c]) + } + } + } + if got[a] != 150 || got[b] != 50 || got[c] != 200 { + t.Fatalf("over 400: a %d b %d c %d", got[a], got[b], got[c]) + } + if run := strings.Count(strings.Join(seq[:8], " "), c+" "+c+" "+c); run > 0 { + t.Fatalf("not smooth: %v", seq[:8]) + } + + // c resting: a and b share its requests 3 : 1, and c is tried last + restingUntil.Lock() + restingUntil.m[c] = time.Now().Add(time.Hour) + restingUntil.Unlock() + got = map[string]int{} + for range 4 * 25 { + cs := route(p, perKey(p, "m", provider.Chat), "m", provider.Chat) + cs, _ = restLast(cs, planned{order: make([]Weighed, len(cs))}) + if cs[len(cs)-1].rest != c { + t.Fatalf("the resting key isn't last: %s", restsOf(cs)) + } + got[cs[0].rest]++ + } + if got[a] != 75 || got[b] != 25 || got[c] != 0 { + t.Fatalf("with c resting: a %d b %d c %d", got[a], got[b], got[c]) + } + + // a conversation kept on b counts against b: the next picks make it up + restingUntil.Lock() + delete(restingUntil.m, c) + restingUntil.Unlock() + wrr.Lock() + delete(wrr.m, p.ID) + wrr.Unlock() + got = map[string]int{} + for i := range 8 * 10 { + cs := route(p, perKey(p, "m", provider.Chat), "m", provider.Chat) + k := cs[0] + if i%8 == 0 { // every round, one request stays with b + for _, x := range cs { + if x.rest == b { + weightedKept(k, x) + k = x + } + } + } + got[k.rest]++ + } + if got[a] != 30 || got[b] != 10 || got[c] != 40 { + t.Fatalf("with b kept: a %d b %d c %d", got[a], got[b], got[c]) + } +} + +// A key's weight is saved with it, the first key's too, and 1 is the +// default left out of the file. +func TestSetKeyWeight(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + t.Setenv("USERPROFILE", home) + t.Setenv("XDG_CONFIG_HOME", home+"/.config") + if err := provider.Save(provider.Provider{ID: "wk", Name: "WK", Chat: "https://wk.example/v1", Key: "ka", Keys: []provider.KeyAccount{{Key: "kb"}}, Routing: provider.Weighted}); err != nil { + t.Fatal(err) + } + if err := provider.SetKeyWeight("wk", provider.KeyID("ka"), 3); err != nil { + t.Fatal(err) + } + if err := provider.SetKeyWeight("wk", provider.KeyID("kb"), 5); err != nil { + t.Fatal(err) + } + if err := provider.SetKeyWeight("wk", provider.KeyID("kb"), -1); err == nil { + t.Fatal("a weight below 0 was taken") + } + p, err := provider.Find("wk") + if err != nil { + t.Fatal(err) + } + if p.Routing != provider.Weighted { + t.Fatalf("routing %q", p.Routing) + } + ks := p.KeyList() + if ks[0].Weight != 3 || ks[1].Weight != 5 { + t.Fatalf("weights %d %d", ks[0].Weight, ks[1].Weight) + } + // the weight goes with its key when another goes first + if err := provider.UseKey("wk", provider.KeyID("kb")); err != nil { + t.Fatal(err) + } + p, _ = provider.Find("wk") + if p.KeyWeight != 5 || p.Keys[0].Weight != 3 { + t.Fatalf("after Make first: first %d, other %d", p.KeyWeight, p.Keys[0].Weight) + } + if err := provider.SetKeyWeight("wk", provider.KeyID("kb"), 1); err != nil { + t.Fatal(err) + } + p, _ = provider.Find("wk") + if p.KeyWeight != 0 { + t.Fatalf("1 kept as %d", p.KeyWeight) + } +} diff --git a/internal/gui/assets/app.css b/internal/gui/assets/app.css index 5456d14f..8744e81a 100644 --- a/internal/gui/assets/app.css +++ b/internal/gui/assets/app.css @@ -1783,6 +1783,9 @@ body.window .ag-sk-icon { width: 22px; height: 22px; border-radius: 6px; } @media (max-width: 520px) { .acc.adding .kf { grid-template-columns: 1fr; } } .acc.adding .bulk-keys { width: 100%; box-sizing: border-box; resize: vertical; min-height: 84px; font-family: var(--code); font-size: 11.5px; line-height: 1.5; padding: 6px 8px; border: 1px solid var(--line); border-radius: 6px; background: var(--card); color: var(--fg); } .acc.adding .bulk-count { margin-left: 6px; white-space: nowrap; } +.acc .key-weight { display: none; flex: none; font-size: 11px; padding: 1px 6px; border-radius: 5px; color: var(--fg-2); background: var(--pill); border: 0; cursor: pointer; white-space: nowrap; font-variant-numeric: tabular-nums; } +.accts.weighted .acc .key-weight { display: inline-block; } +.acc .key-weight-in { width: 64px; flex: none; font-size: 12px; padding: 1px 6px; } .acc .key-rest { flex: none; font-size: 11px; padding: 1px 6px; border-radius: 5px; color: var(--amber); background: var(--amber-soft); white-space: nowrap; } .acc .proto { flex: none; display: inline-flex; align-items: center; gap: 3px; height: 20px; padding: 0 5px 0 7px; font: inherit; font-size: 11px; color: var(--muted); border: 0; border-radius: 5px; background: var(--pill); white-space: nowrap; cursor: default; transition: opacity .15s, background .12s, color .12s; } .acc .proto svg { opacity: .7; transition: transform .15s var(--ease); } diff --git a/internal/gui/assets/app.js b/internal/gui/assets/app.js index 722b16ca..f5a98300 100644 --- a/internal/gui/assets/app.js +++ b/internal/gui/assets/app.js @@ -8261,7 +8261,10 @@ const ROUTINGS = [ ["rotate", "In turn", "Each turn of a conversation goes to the next one, spreading the load evenly; the requests within a turn stay where it began, so the prompt cache holds, and one that fails is passed over while it rests."], ["usage", "Least used first", "Each request goes to the one used least: a subscription by the share of its allowance used, a key by the tokens it served in the last hours."], ["pace", "Weekly pace", "Each request goes to the subscription with the most of its week left per hour until it renews — the one with the most to lose at its reset — so less of each week is lost at its reset; one at 90% or more waits until the others can't answer. A key goes by the tokens it served in the last hours."], + ["weight", "By weight", "Requests spread over the keys by the weight set beside each: a key weighing 3 takes three requests for every one a key weighing 1 takes, evenly over a few requests. One that fails is passed over while it rests, and the others share its requests; a conversation stays with its key as Stays says."], ]; +// weight is a key's share: offered for a provider with keys only +const routingsOf = (p) => p.account ? ROUTINGS.filter(([id]) => id !== "weight") : ROUTINGS; // as on the Routing page's list of these (routing.js AFF_OPTS) const STAYS = [ ["", "Auto", "A conversation stays with the account or key that answered it while what the vendor cached of it is worth keeping — within a turn always, across turns while it's fresh."], @@ -8315,10 +8318,12 @@ function renderRouting(p) { const cur = ROUTINGS.find(([id]) => id === routingNow()) || ROUTINGS[0]; rHint.textContent = t(cur[2]) + own(cur[0]); rUnsaved.hidden = routingNow() === (p.routing || ""); - // in turn goes round already - kLabel.hidden = kWrap.hidden = routingNow() === "rotate"; + // in turn and by weight go round already + kLabel.hidden = kWrap.hidden = routingNow() === "rotate" || routingNow() === "weight"; + // each key's weight shows beside it while the routing is by weight + for (const l of document.querySelectorAll(".accts.keys")) if (l.dataset.provider === p.id) l.classList.toggle("weighted", routingNow() === "weight"); }; - const rPick = segs(ROUTINGS.map(([id, name]) => [id, t(name)]), routingNow(), (routing) => { draft.routing = routing; drawRouting(); }); + const rPick = segs(routingsOf(p).map(([id, name]) => [id, t(name)]), routingNow(), (routing) => { draft.routing = routing; drawRouting(); }); const rRow = el("div", "route-pick"); rRow.append(rPick, rUnsaved); const rWrap = el("div"); @@ -10013,12 +10018,45 @@ function accountModels(p, ref, isKey, name) { return [pill, box]; } +// keyWeight: a key's weight under By weight routing, and its share of the +// requests; a click edits it in place, Enter or leaving it saves it. +function keyWeight(p, k, w, total) { + const b = el("button", "key-weight", t("Weight {n} · {share}%", { n: w, share: Math.round(100 * w / total) })); + b.title = t("Its share of the requests: a key weighing 3 takes three for every one a key weighing 1 takes"); + b.onclick = () => { + const i = input(String(w), "1"); + i.type = "number"; + i.min = "1"; + i.max = "1000"; + i.className = "key-weight-in"; + let done = false; + const save = (keep) => { + if (done) return; + done = true; + const n = Math.round(Number(i.value)); + if (keep && n >= 1 && n <= 1000 && n !== w) accountAction("keys/weight", { id: p.id, ref: k.id, weight: n }, t("{key} weighs {n}", { key: k.name || k.masked, n })); + else renderProviders(); + }; + i.onkeydown = (e) => { e.stopPropagation(); if (e.key === "Enter") save(true); else if (e.key === "Escape") save(false); }; + i.onblur = () => save(true); + b.replaceWith(i); + i.focus({ preventScroll: true }); + i.select(); + }; + return b; +} + // renderKeyAccounts: a key provider's accounts, one per key, the same list // a subscription has. addingKey holds the half-typed new one. let addingKey = null; function renderKeyAccounts(p) { - const list = el("div", "accts"); + const list = el("div", "accts keys"); + list.dataset.provider = p.id; + list.classList.toggle("weighted", (draft?.routing ?? (p.routing || "")) === "weight"); const several = p.keyList.filter((k) => k.on).length > 1; + // the weights of those on, for each one's share (#841) + const weightOf = (k) => Math.max(k.weight || 0, 1); + const weights = p.keyList.filter((k) => k.on).reduce((n, k) => n + weightOf(k), 0); for (const k of p.keyList) { const row = el("div", "acc" + (k.on ? " in-use" : " off") + (k.id === justAdded ? " new" : "")); row.dataset.accountId = k.id; @@ -10049,6 +10087,7 @@ function renderKeyAccounts(p) { r.title = t("The gateway passes this key over until then, and tries the next one"); row.append(r); } + if (k.on && several) row.append(keyWeight(p, k, weightOf(k), weights)); const proto = protoPicker(p, k.protocol, (v) => accountAction("keys/protocol", { id: p.id, ref: k.id, protocol: v })); if (proto) row.append(proto); // its own models (#474), when there is another key to send the rest to diff --git a/internal/gui/assets/i18n.js b/internal/gui/assets/i18n.js index 2c5b8df1..05847ff6 100644 --- a/internal/gui/assets/i18n.js +++ b/internal/gui/assets/i18n.js @@ -790,6 +790,12 @@ const I18N = { "In order": "按顺序", "In turn": "轮流", "Least used first": "用量少的优先", + "By weight": "按权重", + "By weight: it's {who}'s share \u2014 each key takes requests as its weight says.": "按权重：轮到 {who} 的份额——每个密钥按自己的权重接请求。", + "Requests spread over the keys by the weight set beside each: a key weighing 3 takes three requests for every one a key weighing 1 takes, evenly over a few requests. One that fails is passed over while it rests, and the others share its requests; a conversation stays with its key as Stays says.": "请求按每个密钥旁设置的权重分配：权重 3 的密钥每接 3 个请求，权重 1 的接 1 个，几个请求之内就均匀分开。出错的在休息期间被跳过，由其他密钥按各自权重分担；会话按「会话保持」的设置留在原来的密钥上。", + "Weight {n} · {share}%": "权重 {n} · {share}%", + "Its share of the requests: a key weighing 3 takes three for every one a key weighing 1 takes": "它在请求中的份额：权重 3 的密钥每接 3 个请求，权重 1 的接 1 个", + "{key} weighs {n}": "{key} 的权重已设为 {n}", "Requests go to the first; the next takes over when the one before runs out of quota, hits a rate limit or fails.": "请求先发给第一个；它额度用完、被限流或出错时，由下一个接手。", "Each turn of a conversation goes to the next one, spreading the load evenly; the requests within a turn stay where it began, so the prompt cache holds, and one that fails is passed over while it rests.": "会话的每一轮轮流交给下一个，均摊负载；一轮之内的请求留在这一轮开始的地方，好让提示词缓存生效；出错的在休息期间被跳过。", "Each request goes to the one used least: a subscription by the share of its allowance used, a key by the tokens it served in the last hours.": "每个请求发给用量最少的：订阅看已用额度比例，密钥看最近几小时经它的 token 数。", @@ -3840,6 +3846,12 @@ const I18N = { "In order": "順番", "In turn": "交互", "Least used first": "使用量が最少のものを優先", + "By weight": "重み付け", + "By weight: it's {who}'s share \u2014 each key takes requests as its weight says.": "重み付け：{who} の取り分です——各キーは重みどおりにリクエストを受けます。", + "Requests spread over the keys by the weight set beside each: a key weighing 3 takes three requests for every one a key weighing 1 takes, evenly over a few requests. One that fails is passed over while it rests, and the others share its requests; a conversation stays with its key as Stays says.": "各キーの横で設定した重みでリクエストを振り分けます。重み 3 のキーは、重み 1 のキーが 1 件受けるごとに 3 件受け、数件のうちに均等に分かれます。失敗したキーは休止中は飛ばされ、他のキーがそれぞれの重みで分担します。会話は「維持」の設定どおりそのキーにとどまります。", + "Weight {n} · {share}%": "重み {n} · {share}%", + "Its share of the requests: a key weighing 3 takes three for every one a key weighing 1 takes": "リクエストの取り分：重み 3 のキーは、重み 1 のキーが 1 件受けるごとに 3 件受けます", + "{key} weighs {n}": "{key} の重みを {n} にしました", "Requests go to the first; the next takes over when the one before runs out of quota, hits a rate limit or fails.": "リクエストは最初のものに送られます。クォータ切れ、レート制限、エラーのときは次のものが引き継ぎます。", "Each turn of a conversation goes to the next one, spreading the load evenly; the requests within a turn stay where it began, so the prompt cache holds, and one that fails is passed over while it rests.": "会話のターンごとに次のものへ送り、負荷を均等に分散します。ターン内のリクエストは開始した場所に留まるのでプロンプトキャッシュが効きます。エラーになったものは休止中スキップされます。", "Each request goes to the one used least: a subscription by the share of its allowance used, a key by the tokens it served in the last hours.": "各リクエストは使用量が最も少ないものに送られます。サブスクリプションは利用枠の使用割合、キーは直近数時間に処理したトークン数で判断します。", @@ -6882,6 +6894,12 @@ const I18N = { "In order": "Der Reihe nach", "In turn": "Abwechselnd", "Least used first": "Am wenigsten genutzte zuerst", + "By weight": "Nach Gewicht", + "By weight: it's {who}'s share \u2014 each key takes requests as its weight says.": "Nach Gewicht: {who} ist mit seinem Anteil dran — jeder Schlüssel bekommt Anfragen nach seinem Gewicht.", + "Requests spread over the keys by the weight set beside each: a key weighing 3 takes three requests for every one a key weighing 1 takes, evenly over a few requests. One that fails is passed over while it rests, and the others share its requests; a conversation stays with its key as Stays says.": "Anfragen verteilen sich nach dem Gewicht neben jedem Schlüssel: Einer mit Gewicht 3 bekommt drei Anfragen für jede, die einer mit Gewicht 1 bekommt, schon über wenige Anfragen gleichmäßig. Ein fehlschlagender wird übersprungen, solange er pausiert, und die anderen teilen sich seine Anfragen; ein Gespräch bleibt bei seinem Schlüssel, wie „Bleibt“ es sagt.", + "Weight {n} · {share}%": "Gewicht {n} · {share}%", + "Its share of the requests: a key weighing 3 takes three for every one a key weighing 1 takes": "Sein Anteil an den Anfragen: Ein Schlüssel mit Gewicht 3 bekommt drei für jede, die einer mit Gewicht 1 bekommt", + "{key} weighs {n}": "{key} hat jetzt Gewicht {n}", "Requests go to the first; the next takes over when the one before runs out of quota, hits a rate limit or fails.": "Anfragen gehen an das erste; das nächste übernimmt, wenn das vorherige kein Kontingent mehr hat, ein Rate-Limit erreicht oder fehlschlägt.", "Each turn of a conversation goes to the next one, spreading the load evenly; the requests within a turn stay where it began, so the prompt cache holds, and one that fails is passed over while it rests.": "Jede Runde eines Gesprächs geht an das nächste, damit sich die Last gleichmäßig verteilt; die Anfragen einer Runde bleiben dort, wo sie begann, damit der Prompt-Cache greift, und ein fehlgeschlagenes wird übersprungen, solange es pausiert.", "Each request goes to the one used least: a subscription by the share of its allowance used, a key by the tokens it served in the last hours.": "Jede Anfrage geht an das am wenigsten genutzte: ein Abo nach dem verbrauchten Anteil seines Kontingents, ein Schlüssel nach den Token, die er in den letzten Stunden geliefert hat.", diff --git a/internal/gui/assets/routing.js b/internal/gui/assets/routing.js index fbb71482..590b2b23 100644 --- a/internal/gui/assets/routing.js +++ b/internal/gui/assets/routing.js @@ -297,6 +297,7 @@ ? t("In order: with {rested} resting after a failure, {who} is the first that can answer.", { rested: rested.join(", "), who: w }) : t("In order: {who} is first, and answers everything while it can.", { who: w }); case "rotate": return t("In turn: it's {who}'s turn — each request starts one further along.", { who: w }); + case "weight": return t("By weight: it's {who}'s share — each key takes requests as its weight says.", { who: w }); case "usage": if (f.kind === "account" && f.known) return t("Least used first: {who} has the most of its allowance left — {n} used.", { who: w, n: pct(f.used) }); if (f.kind === "key") return t("Least used first: {who} served the fewest tokens lately — {n}.", { who: w, n: tokens(f.tokens || 0) }); @@ -2291,7 +2292,7 @@ // in turn goes round already, and a manual group sends to one member // how long a group's member may take to its first token (provider.Group.FirstToken) const FIRST_OPTS = [[0, "Wait"], [30, "30 s"], [60, "1 min"], [120, "2 min"]]; - const sinkable = (routing) => routing !== "rotate" && routing !== "manual"; + const sinkable = (routing) => routing !== "rotate" && routing !== "weight" && routing !== "manual"; const AFF_HINT = { "": "A conversation stays with the account or key that answered it while what the vendor cached of it is worth keeping — within a turn always, across turns while it's fresh.", session: "A conversation stays with the account or key that answered it for the whole session, while it can answer.", @@ -3286,8 +3287,10 @@ try { await api("provider/" + what, body); status(ok, "ok"); groups = await api("groups"); load(); } catch (e) { status(e.message, "err"); } }; const lab = (s, c) => { const w = el("span", "lab"); w.append(el("small", "", s), c); return w; }; + // by weight is a key's share (#841), set beside each key in the provider's editor + const opts = p.kind === "account" ? ROUTE_OPTS : [...ROUTE_OPTS, ["weight", "By weight"]]; ctl.append( - lab(t("Routing"), segs(ROUTE_OPTS.map(([id, n]) => [id, t(n)]), p.routing || "", (v) => set("route", { id: p.provider, routing: v }, t("{name}: {routing}", { name: p.name, routing: t(ROUTE_OPTS.find(([id]) => id === v)[1]) })))), + lab(t("Routing"), segs(opts.map(([id, n]) => [id, t(n)]), p.routing || "", (v) => set("route", { id: p.provider, routing: v }, t("{name}: {routing}", { name: p.name, routing: t(opts.find(([id]) => id === v)[1]) })))), lab(t("Stays"), segs(AFF_OPTS.map(([id, n]) => [id, t(n)]), p.affinity || "", (v) => set("affinity", { id: p.provider, affinity: v }, t("{name}: {routing}", { name: p.name, routing: t(AFF_OPTS.find(([id]) => id === v)[1]) }))))); if (sinkable(p.routing || "")) { const k = lab(t("Rate limited"), segs(SINK_OPTS.map(([id, n]) => [id, t(n)]), p.sink ? "sink" : "", (v) => set("sink", { id: p.provider, sink: v === "sink" }, t("{name}: {routing}", { name: p.name, routing: t(SINK_OPTS.find(([id]) => id === v)[1]) })))); diff --git a/internal/gui/providers.go b/internal/gui/providers.go index d1292064..558239c2 100644 --- a/internal/gui/providers.go +++ b/internal/gui/providers.go @@ -987,7 +987,7 @@ func providerRoutes(mux *http.ServeMux, w Windows) { in.Contexts = old.Contexts // a save that doesn't say } if in.Key == old.Key { - in.KeyName, in.KeyProtocol = old.KeyName, old.KeyProtocol + in.KeyName, in.KeyProtocol, in.KeyWeight = old.KeyName, old.KeyProtocol, old.KeyWeight } } if in.Icon == "" && old != nil && in.Preset == "" { @@ -1358,6 +1358,7 @@ func providerRoutes(mux *http.ServeMux, w Windows) { var in struct { ID, Key, Name, Ref string Protocol provider.Protocol + Weight int } var added, had int if err := json.NewDecoder(r.Body).Decode(&in); err != nil { @@ -1374,6 +1375,8 @@ func providerRoutes(mux *http.ServeMux, w Windows) { added, had, err = provider.AddKeys(in.ID, provider.SplitKeys(in.Key), in.Protocol) case "protocol": err = provider.SetKeyProtocol(in.ID, in.Ref, in.Protocol) + case "weight": + err = provider.SetKeyWeight(in.ID, in.Ref, in.Weight) case "use": err = provider.UseKey(in.ID, in.Ref) case "remove": @@ -1543,7 +1546,7 @@ func providerRoutes(mux *http.ServeMux, w Windows) { // all) is the saved one's; the provider's other keys are kept. func typed(p, in provider.Provider, proxy *string) provider.Provider { if k := strings.TrimSpace(in.Key); k != "" && k != p.Key { - p.Key, p.KeyName, p.KeyProtocol = k, "", "" + p.Key, p.KeyName, p.KeyProtocol, p.KeyWeight = k, "", "", 0 } for _, f := range []struct { to *string diff --git a/internal/gui/tests/key-weight.test.cjs b/internal/gui/tests/key-weight.test.cjs new file mode 100644 index 00000000..487657e0 --- /dev/null +++ b/internal/gui/tests/key-weight.test.cjs @@ -0,0 +1,114 @@ +// Run with Node's test runner and Playwright on the module path; see README.md. +// Weighted keys (#841): a key provider's Routing has By weight, and with it +// picked each key that's on shows its weight and its share of the +// requests beside it; a click edits the weight in place and posts +// keys/weight. Hidden under the other routings, and offered to key +// providers only. The page doesn't scroll. In English and Chinese. +const assert = require("node:assert/strict"); +const fs = require("node:fs/promises"); +const path = require("node:path"); +const { test } = require("node:test"); +const { chromium, webkit } = require("playwright"); + +const assets = path.resolve(__dirname, "../assets"); +const relay = () => ({ + id: "relay", name: "Relay", icon: "generic", chat: "https://relay.example.com/v1", responses: "", anthropic: "", catalog: "", + models: [{ id: "big", name: "", on: true }], agents: [], fallback: [], headers: {}, routing: "", + key: { set: true, masked: "sk-…one" }, balanceToken: { takes: false, set: false }, proxy: "", + keyList: [ + { id: "aaaaaaaaaa", name: "team-a", masked: "sk-…one", on: true, active: true, weight: 3 }, + { id: "bbbbbbbbbb", name: "team-b", masked: "sk-…two", on: true }, + { id: "cccccccccc", name: "", masked: "sk-…off", on: false }, + ], +}); + +function serve(lang, posts) { + const p = relay(); + const providers = { providers: [p], presets: [], excluded: [], gateway: { running: true, window: true } }; + const state = { agents: [], profiles: [], settings: { lang, theme: "light" } }; + return async (route) => { + const url = new URL(route.request().url()); + const json = (data) => route.fulfill({ json: data }); + if (url.pathname === "/boot.js") return route.fulfill({ contentType: "text/javascript", body: `window.bootPrefs = {lang:"${lang}",theme:"light",web:true};` }); + if (url.pathname === "/wails/runtime.js") return route.fulfill({ contentType: "text/javascript", body: "export const Window = {};" }); + if (url.pathname === "/api/state") return json(state); + if (url.pathname === "/api/providers") return json(providers); + if (url.pathname === "/api/groups") return json({ groups: [] }); + if (url.pathname === "/api/plugins") return json({ plugins: [] }); + if (url.pathname === "/api/keys/weight" && route.request().method() === "POST") { + const body = route.request().postDataJSON(); + posts.push(body); + p.keyList = p.keyList.map((k) => k.id === body.ref ? { ...k, weight: body.weight } : k); + return json(providers); + } + if (url.pathname.startsWith("/api/")) return json({}); + const file = path.join(assets, url.pathname === "/" ? "index.html" : url.pathname); + const contentType = { ".html": "text/html", ".js": "text/javascript", ".css": "text/css", ".svg": "image/svg+xml", ".png": "image/png" }[path.extname(file)]; + await route.fulfill({ body: await fs.readFile(file), contentType }); + }; +} + +const words = { + en: { by: "By weight", a: "Weight 3 · 75%", b: "Weight 1 · 25%", b2: "Weight 2 · 40%", a2: "Weight 3 · 60%", note: "team-b weighs 2" }, + zh: { by: "按权重", a: "权重 3 · 75%", b: "权重 1 · 25%", b2: "权重 2 · 40%", a2: "权重 3 · 60%", note: "team-b 的权重已设为 2" }, +}; + +for (const engine of (process.env.BROWSER ? [process.env.BROWSER] : ["chromium", "webkit"])) { + for (const lang of ["en", "zh"]) { + const w = words[lang]; + test(`${engine} ${lang}: By weight shows each key's weight and share, and a click sets it`, async (t) => { + const browser = await (engine === "webkit" ? webkit.launch() : chromium.launch({ channel: "chromium" })); + const page = await (await browser.newContext({ viewport: { width: 900, height: 760 }, reducedMotion: "reduce" })).newPage(); + t.after(async () => { + if (process.env.ARTIFACT_DIR) { + await fs.mkdir(process.env.ARTIFACT_DIR, { recursive: true }); + await page.screenshot({ path: path.join(process.env.ARTIFACT_DIR, `${engine}-${lang}-key-weight.png`) }); + } + await browser.close(); + }); + page.setDefaultTimeout(5000); + const errors = []; + page.on("pageerror", (e) => errors.push(e.message)); + const posts = []; + await page.route("**/*", serve(lang, posts)); + await page.goto("http://magpie.test/?view=providers"); + await page.locator(".row.provider", { hasText: "Relay" }).first().click(); + await page.locator(".editor .accts .acc.add").waitFor(); + const scrolled = () => page.evaluate(() => [window.scrollY, ...[...document.querySelectorAll("*")].filter((e) => e.scrollTop > 0).map((e) => `${e.className}:${e.scrollTop}`)].join(" ")); + const chip = (id) => page.locator(`.acc[data-account-id="${id}"] .key-weight`); + + // smart: no weights shown + assert.equal(await chip("aaaaaaaaaa").isVisible(), false); + const by = page.locator(".editor").getByRole("button", { name: w.by, exact: true }); + await by.scrollIntoViewIfNeeded(); + const sc = await scrolled(); + await by.click(); + assert.equal(await scrolled(), sc, "a click scrolled"); + assert.equal((await chip("aaaaaaaaaa").textContent()).trim(), w.a); + assert.equal((await chip("bbbbbbbbbb").textContent()).trim(), w.b); + assert.equal(await chip("aaaaaaaaaa").isVisible(), true); + assert.equal(await chip("cccccccccc").count(), 0, "a key that's off has no share"); + + // a click edits it in place, Enter saves it + await chip("bbbbbbbbbb").click(); + const box = page.locator('.acc[data-account-id="bbbbbbbbbb"] .key-weight-in'); + await box.waitFor(); + assert.equal(await scrolled(), sc, "a click scrolled"); + await box.fill("2"); + await box.press("Enter"); + for (let i = 0; i < 60 && !posts.length; i++) await new Promise((r) => setTimeout(r, 50)); + assert.deepEqual(posts.at(-1), { id: "relay", ref: "bbbbbbbbbb", weight: 2 }); + await page.locator(`.acc[data-account-id="bbbbbbbbbb"] .key-weight`, { hasText: w.b2 }).waitFor(); + assert.equal((await chip("aaaaaaaaaa").textContent()).trim(), w.a2); + // still shown: the routing picked stays in the editor's draft + assert.equal(await chip("bbbbbbbbbb").isVisible(), true); + assert.equal((await page.locator("#status").textContent()).trim(), w.note); + + if (process.env.ARTIFACT_DIR) await page.locator(".editor .accts").screenshot({ path: path.join(process.env.ARTIFACT_DIR, `${engine}-${lang}-key-weight-rows.png`) }); + // back to Smart hides them + await page.locator(".editor").getByRole("button", { name: lang === "en" ? "Smart" : "智能", exact: true }).first().click(); + assert.equal(await chip("aaaaaaaaaa").isVisible(), false); + assert.deepEqual(errors, []); + }); + } +} diff --git a/internal/provider/backup.go b/internal/provider/backup.go index f2396677..769da734 100644 --- a/internal/provider/backup.go +++ b/internal/provider/backup.go @@ -40,7 +40,7 @@ func Restore(ps []Provider) (added, replaced int, err error) { continue } if p.Key == "" && len(p.Keys) == 0 { - p.Key, p.KeyName, p.Keys, p.KeyProtocol = f.Providers[i].Key, f.Providers[i].KeyName, f.Providers[i].Keys, f.Providers[i].KeyProtocol + p.Key, p.KeyName, p.Keys, p.KeyProtocol, p.KeyWeight = f.Providers[i].Key, f.Providers[i].KeyName, f.Providers[i].Keys, f.Providers[i].KeyProtocol, f.Providers[i].KeyWeight if p.BalanceToken == "" { p.BalanceToken = f.Providers[i].BalanceToken } @@ -104,7 +104,7 @@ func Mirror(ps []Provider, gs []Group) error { continue } if h, ok := here[p.ID]; ok && p.Key == "" && len(p.Keys) == 0 { - p.Key, p.KeyName, p.Keys, p.KeyProtocol = h.Key, h.KeyName, h.Keys, h.KeyProtocol + p.Key, p.KeyName, p.Keys, p.KeyProtocol, p.KeyWeight = h.Key, h.KeyName, h.Keys, h.KeyProtocol, h.KeyWeight if p.BalanceToken == "" { p.BalanceToken = h.BalanceToken } diff --git a/internal/provider/keys.go b/internal/provider/keys.go index 576b2e03..b3af5d01 100644 --- a/internal/provider/keys.go +++ b/internal/provider/keys.go @@ -24,6 +24,10 @@ type KeyAccount struct { // Protocol, when set, is the only one the key works with: some relays // give out one key for Anthropic and another for OpenAI. Empty is any. Protocol Protocol `json:"protocol,omitempty"` + // Weight is the key's share of the requests when the provider's + // routing is Weighted (#841): one with 3 takes three for every one a + // key with 1 takes. None, or 0, counts as 1. + Weight int `json:"weight,omitempty"` } // KeyInfo describes one of a provider's keys without giving it away. @@ -35,6 +39,7 @@ type KeyInfo struct { On bool `json:"on"` // in use: the first, or next in line Protocol Protocol `json:"protocol,omitempty"` // the only one it works with + Weight int `json:"weight,omitempty"` // its share when routing is Weighted // Rest is why the gateway passes it over now, after a failure, and // until when; nil while it takes requests. @@ -71,10 +76,10 @@ func KeyID(key string) string { return keyID(key) } func (p Provider) KeyList() []KeyInfo { var out []KeyInfo if p.Key != "" { - out = append(out, KeyInfo{ID: keyID(p.Key), Name: p.KeyName, Masked: Mask(p.Key), Active: true, On: true, Protocol: p.KeyProtocol}) + out = append(out, KeyInfo{ID: keyID(p.Key), Name: p.KeyName, Masked: Mask(p.Key), Active: true, On: true, Protocol: p.KeyProtocol, Weight: p.KeyWeight}) } for _, k := range p.Keys { - out = append(out, KeyInfo{ID: keyID(k.Key), Name: k.Name, Masked: Mask(k.Key), On: !k.Off, Protocol: k.Protocol}) + out = append(out, KeyInfo{ID: keyID(k.Key), Name: k.Name, Masked: Mask(k.Key), On: !k.Off, Protocol: k.Protocol, Weight: k.Weight}) } return out } @@ -187,11 +192,11 @@ func AddKeys(id string, keys []string, proto Protocol) (added, had int, err erro // first is the first key as a KeyAccount, and setFirst makes k the first. func (p *Provider) first() KeyAccount { - return KeyAccount{Name: p.KeyName, Key: p.Key, Protocol: p.KeyProtocol} + return KeyAccount{Name: p.KeyName, Key: p.Key, Protocol: p.KeyProtocol, Weight: p.KeyWeight} } func (p *Provider) setFirst(k KeyAccount) { - p.Key, p.KeyName, p.KeyProtocol = k.Key, k.Name, k.Protocol + p.Key, p.KeyName, p.KeyProtocol, p.KeyWeight = k.Key, k.Name, k.Protocol, k.Weight } func keyProtocolOK(proto Protocol) error { @@ -223,10 +228,41 @@ func SetKeyProtocol(id, keyRef string, proto Protocol) error { return Save(*p) } +// MaxKeyWeight is the most a key's weight can be. +const MaxKeyWeight = 1000 + +// SetKeyWeight sets a key's share of the requests under Weighted routing; +// 0 takes it back to the default, 1. +func SetKeyWeight(id, keyRef string, weight int) error { + if weight < 0 || weight > MaxKeyWeight { + return fmt.Errorf("a key's weight is 0 to %d", MaxKeyWeight) + } + p, err := Find(id) + if err != nil { + return err + } + i, ok := findKey(p, keyRef) + if !ok { + return fmt.Errorf("%s has no such key", p.Name) + } + if weight == 1 { + weight = 0 // the default: kept out of the file + } + if i < 0 { + p.KeyWeight = weight + } else { + p.Keys[i].Weight = weight + } + return Save(*p) +} + +// WeightOf is the key's weight as Weighted routing counts it: 1 when none. +func (k KeyAccount) WeightOf() int { return max(k.Weight, 1) } + // WithKey is p using key k: its endpoints narrowed to k's protocol when k // has one. It has none left when p doesn't serve that protocol. func (p Provider) WithKey(k KeyAccount) Provider { - p.Key, p.KeyName, p.KeyProtocol = k.Key, k.Name, k.Protocol + p.Key, p.KeyName, p.KeyProtocol, p.KeyWeight = k.Key, k.Name, k.Protocol, k.Weight if k.Protocol != "" { for _, pr := range Protocols { if pr != k.Protocol { diff --git a/internal/provider/provider.go b/internal/provider/provider.go index b928e1bd..0ddce0e0 100644 --- a/internal/provider/provider.go +++ b/internal/provider/provider.go @@ -59,6 +59,8 @@ type Provider struct { // for: a relay that hands out one key for Anthropic and another for // OpenAI (see KeyAccount.Protocol). KeyProtocol Protocol `json:"keyProtocol,omitempty"` + // KeyWeight is the first key's weight (see KeyAccount.Weight). + KeyWeight int `json:"keyWeight,omitempty"` // Base URLs, one per protocol the vendor serves natively. magpie appends // the usual paths: chat/responses bases end in /v1 (OpenAI style), @@ -83,7 +85,8 @@ type Provider struct { // can't take it; "rotate" each in turn; "usage" the least used first; // "pace" the one with the most remaining allowance per hour until // its window resets first, so less allowance is lost at reset. - // The window can be shorter than a week. + // The window can be shorter than a week. "weight" by each key's + // weight, a key with 3 taking three requests to one with 1's. // Whichever it is, one out of credit, out of quota, rate limited or // failing is passed over for as long as that lasts. Routing string `json:"routing,omitempty"` @@ -583,7 +586,7 @@ func AddCopy(p Provider, from string) (string, error) { return "", fmt.Errorf("%s is a signed-in account, which can't be copied", src.Name) } if p.Key == "" { - p.Key, p.KeyName, p.KeyProtocol = src.Key, src.KeyName, src.KeyProtocol + p.Key, p.KeyName, p.KeyProtocol, p.KeyWeight = src.Key, src.KeyName, src.KeyProtocol, src.KeyWeight p.Keys = slices.Clone(src.Keys) p.Routing, p.Sink, p.Affinity = src.Routing, src.Sink, src.Affinity } @@ -835,7 +838,7 @@ func normalize(p Provider) Provider { p.Decide = pr.Decide } } - if p.Routing != Ordered && p.Routing != Rotate && p.Routing != LeastUsed && p.Routing != Pace { + if p.Routing != Ordered && p.Routing != Rotate && p.Routing != LeastUsed && p.Routing != Pace && p.Routing != Weighted { p.Routing = "" } if !slices.Contains(Affinities, p.Affinity) { diff --git a/internal/provider/routing.go b/internal/provider/routing.go index 3982605b..f8f734c3 100644 --- a/internal/provider/routing.go +++ b/internal/provider/routing.go @@ -24,6 +24,9 @@ const ( Rotate = "rotate" LeastUsed = "usage" Pace = "pace" + // Weighted spreads a provider's requests over its keys by each key's + // weight (KeyAccount.Weight, #841), smoothly: 3 and 1 go a, a, b, a… + Weighted = "weight" ) // SetRouting changes how a provider's requests spread over its keys or diff --git a/internal/tui/routing.go b/internal/tui/routing.go index 2b1013ff..59b8f0fe 100644 --- a/internal/tui/routing.go +++ b/internal/tui/routing.go @@ -46,6 +46,8 @@ func routingName(v string) string { return "least used" case provider.Pace: return "weekly pace" + case provider.Weighted: + return "by weight" case provider.Manual: return "manual" }