jianf/ModelRouter
1
0
Fork 0
mirror of https://gitcode.com/JianFeeeee/ModelRouter.git synced 2026-09-20 00:48:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(auto): AUTO-only priority chain with tiered slots + live source probing; CSV export w/ key names; audit persistence; fix prompt token accounting & deepseek thinking

Browse Source
This commit is contained in:
root
2026-08-10 00:10:19 +08:00
parent 859d310ad3
commit d48b993010
10 changed files with 663 additions and 210 deletions
Download Patch File Download Diff File Expand all files Collapse all files

107
internal/gateway/chat.go
View File

@ -108,22 +108,37 @@ func (g *Gateway) resolveCands(ctx context.Context, req *chatRequest) ([]*provid
// filterCandsByModels keeps only providers exposing at least one model of the
// scope (used for user keys with a restricted model scope). An "AUTO" scope
// entry means the key is allowed to use any model.
// entry means the key is allowed to use any model. Scope entries with a
// Source pinned to a specific upstream narrow the candidates to that source
// for the matching model.
func filterCandsByModels(cands []*provider.Provider, allow []config.ModelScope) []*provider.Provider {
allowed := make(map[string]bool, len(allow))
for _, m := range allow {
if m.Model == "" || strings.EqualFold(m.Model, "AUTO") {
return cands
}
}
allowed := make(map[string]bool, len(allow))
byModelSrc := map[string]map[string]bool{}
for _, m := range allow {
allowed[m.Model] = true
if m.Source != "" {
if byModelSrc[m.Model] == nil {
byModelSrc[m.Model] = map[string]bool{}
}
byModelSrc[m.Model][m.Source] = true
}
}
out := make([]*provider.Provider, 0, len(cands))
for _, p := range cands {
for _, id := range p.Models() {
if allowed[id] {
out = append(out, p)
break
if !allowed[id] {
continue
}
if srcs := byModelSrc[id]; len(srcs) > 0 && !srcs[p.Name()] {
continue
}
out = append(out, p)
break
}
}
return out
@ -197,7 +212,7 @@ func (g *Gateway) scopeTokens(ctx context.Context, sc config.ModelScope) int64 {
return g.stats.KeyTokens(k)
}
win := AutoPeriodSeconds(sc.Period, sc.Hours)
return g.stats.WindowTokens(sc.Model, win)
return g.stats.WindowTokens(sc.Model, sc.Source, win)
}
func hasScopeModel(list []config.ModelScope, s string) bool {
@ -263,7 +278,7 @@ func (g *Gateway) handleChat(w http.ResponseWriter, r *http.Request) {
if model == "" {
model = g.core.DefaultModel()
}
if strings.EqualFold(strings.TrimSpace(req.Model), "AUTO") {
if isAuto(model) {
plans := g.autoPlans()
if len(plans) == 0 {
writeError(w, http.StatusServiceUnavailable, "no_provider", "no auto slot available (quota exhausted or none configured)")
@ -424,6 +439,39 @@ func effectiveImageModel(model string, cands []*provider.Provider) string {
return model
}
func estimatePromptTokens(req *types.ChatRequest) int64 {
if req == nil {
return 0
}
b, _ := json.Marshal(struct {
Messages []types.ChatMessage `json:"messages"`
Tools []interface{} `json:"tools,omitempty"`
}{Messages: req.Messages, Tools: req.Tools})
if len(b) == 0 {
return 0
}
return int64(len(b)/3 + 1)
}
func estimateTextTokens(parts ...interface{}) int64 {
var n int
for _, p := range parts {
switch v := p.(type) {
case string:
n += len(v)
case json.RawMessage:
n += len(v)
case []types.ToolCall:
b, _ := json.Marshal(v)
n += len(b)
}
}
if n == 0 {
return 0
}
return int64(n/3 + 1)
}
func (g *Gateway) singleChat(w http.ResponseWriter, ctx context.Context, cands []*provider.Provider, req *types.ChatRequest, effective string, rec *Req) {
rec.LatMs = 0
t0 := time.Now()
@ -440,7 +488,13 @@ func (g *Gateway) singleChat(w http.ResponseWriter, ctx context.Context, cands [
rec.OK = true
rec.Status = http.StatusOK
rec.Prompt = int64(resp.TokenUsage.Prompt)
if rec.Prompt == 0 {
rec.Prompt = estimatePromptTokens(req)
}
rec.Compl = int64(resp.TokenUsage.Completion)
if rec.Compl == 0 {
rec.Compl = estimateTextTokens(resp.Content, resp.ReasoningContent, resp.ToolCalls)
}
rec.Source = usedSrc
rec.Model = usedModel
g.writeRec(rec)
@ -495,6 +549,7 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [
if usedModel != "" {
rec.Model = usedModel
}
rec.Prompt = estimatePromptTokens(req)
w.Header().Set("Content-Type", "text/event-stream")
w.Header().Set("Cache-Control", "no-cache")
w.Header().Set("Connection", "keep-alive")
@ -540,7 +595,7 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [
choice.FinishReason = &stop
}
chunk.Choices = []ChunkChoice{choice}
rec.Compl += int64(len(ck.Content)+len(ck.ReasoningContent)) / 3
rec.Compl += int64(len(ck.Content)+len(ck.ReasoningContent)+len(ck.ToolCalls)) / 3
if !send(chunk) {
return
}
@ -575,7 +630,7 @@ func (g *Gateway) autoPlans() []autoPlan {
}
plans := make([]autoPlan, 0, len(rules))
for _, e := range rules {
p := g.core.ProviderForModel(e.Model)
p := g.core.ProviderForSlot(e.Model, e.Source)
if p == nil {
continue
}
@ -583,7 +638,7 @@ func (g *Gateway) autoPlans() []autoPlan {
continue
}
win := AutoPeriodSeconds(e.Period, e.Hours)
if e.TokenQuota > 0 && g.stats.WindowTokens(e.Model, win) >= e.TokenQuota {
if e.TokenQuota > 0 && g.stats.WindowTokens(e.Model, e.Source, win) >= e.TokenQuota {
continue
}
plans = append(plans, autoPlan{p: p, model: e.Model, quota: e.TokenQuota, win: win})
@ -599,12 +654,14 @@ func (g *Gateway) singleChatAuto(w http.ResponseWriter, ctx context.Context, pla
rec.LatMs = 0
t0 := time.Now()
var lastErr error
var lastSrc, lastModel string
for _, pl := range plans {
if !pl.p.Available() {
continue
}
r := *req
r.Model = pl.model
lastSrc, lastModel = pl.p.Name(), pl.model
resp, usedSrc, usedModel, err := g.core.Scheduler().Chat(ctx, scheduler.FromRegistry([]*provider.Provider{pl.p}), &r)
if err != nil {
lastErr = err
@ -614,7 +671,13 @@ func (g *Gateway) singleChatAuto(w http.ResponseWriter, ctx context.Context, pla
rec.OK = true
rec.Status = http.StatusOK
rec.Prompt = int64(resp.TokenUsage.Prompt)
if rec.Prompt == 0 {
rec.Prompt = estimatePromptTokens(&r)
}
rec.Compl = int64(resp.TokenUsage.Completion)
if rec.Compl == 0 {
rec.Compl = estimateTextTokens(resp.Content, resp.ReasoningContent, resp.ToolCalls)
}
rec.Source = usedSrc
rec.Model = usedModel
g.writeRec(rec)
@ -645,6 +708,12 @@ func (g *Gateway) singleChatAuto(w http.ResponseWriter, ctx context.Context, pla
rec.OK = false
rec.Status = http.StatusBadGateway
rec.Err = lastErr.Error()
if rec.Model == "" {
rec.Model = lastModel
}
if rec.Source == "" {
rec.Source = lastSrc
}
g.writeRec(rec)
writeError(w, http.StatusBadGateway, "upstream_error", lastErr.Error())
}
@ -687,12 +756,14 @@ func (g *Gateway) streamChatAuto(w http.ResponseWriter, ctx context.Context, pla
return
}
var lastErr error
var lastSrc, lastModel string
for _, pl := range plans {
if !pl.p.Available() {
continue
}
r := *req
r.Model = pl.model
lastSrc, lastModel = pl.p.Name(), pl.model
chunks, usedSrc, usedModel, err := g.core.Scheduler().ChatStream(ctx, scheduler.FromRegistry([]*provider.Provider{pl.p}), &r)
if err != nil {
lastErr = err
@ -702,6 +773,7 @@ func (g *Gateway) streamChatAuto(w http.ResponseWriter, ctx context.Context, pla
rec.Model = usedModel
rec.Source = usedSrc
}
rec.Prompt = estimatePromptTokens(&r)
for ck := range chunks {
chunk := ChatChunk{
ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model,
@ -719,7 +791,7 @@ func (g *Gateway) streamChatAuto(w http.ResponseWriter, ctx context.Context, pla
choice.FinishReason = &stop
}
chunk.Choices = []ChunkChoice{choice}
rec.Compl += int64(len(ck.Content)+len(ck.ReasoningContent)) / 3
rec.Compl += int64(len(ck.Content)+len(ck.ReasoningContent)+len(ck.ToolCalls)) / 3
if !send(chunk) {
return
}
@ -732,11 +804,14 @@ func (g *Gateway) streamChatAuto(w http.ResponseWriter, ctx context.Context, pla
rec.OK = false
rec.Status = http.StatusBadGateway
rec.Err = lastErr.Error()
stop := "stop"
send(ChatChunk{
ID: id, Object: "chat.completion.chunk", Created: created, Model: "auto",
Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{}, FinishReason: &stop}},
})
if rec.Model == "" {
rec.Model = lastModel
}
if rec.Source == "" {
rec.Source = lastSrc
}
errEvent, _ := json.Marshal(map[string]interface{}{"error": map[string]string{"message": lastErr.Error(), "type": "upstream_error"}})
fmt.Fprintf(w, "data: %s\n\n", errEvent)
fmt.Fprintf(w, "data: [DONE]\n\n")
if flusher != nil {
flusher.Flush()