mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-10-15 12:14:54 +00:00
refactor(gateway): collapse chat.go four-way duplication
singleChat / singleChatAuto / streamChat / streamChatAuto shared ~260 near- identical lines (diff after stripping comments was empty). Extract three shared bodies and shrink all four entry points to thin dispatchers: - failChat: error → record + writeError, shared ChainErr extraction (errors.As returns false for direct-path errors, so the two paths stay equivalent without a branch) - recordChatUsage: exact upstream numbers win, byte estimates fill gaps - writeChatCompletion: unified ChatCompletion rendering (model name is the only direct/auto difference, passed in) - pumpStream: the full SSE pump (preamble, delta loop, terminal finish + usage chunk, [DONE]) — shared by both streaming entries Behavior change (pinned by TestDirectStreamFailoverAuditSource): direct streams now pin rec.Source to the source that actually served the stream after a failover, instead of discarding it (_, usedModel). The audit row previously recorded the first candidate, which was wrong on failover. chat.go: 982 → 896 lines (−86).
This commit is contained in:
@ -702,3 +702,40 @@ type flushRecorder struct {
|
||||
}
|
||||
|
||||
func (f *flushRecorder) Flush() { f.flushed = true }
|
||||
|
||||
// TestDirectStreamFailoverAuditSource: a direct streaming request whose first
|
||||
// candidate hard-fails must be served by the fallback, and the audit record's
|
||||
// Source must name the source that actually served the stream (not just the
|
||||
// first candidate).
|
||||
func TestDirectStreamFailoverAuditSource(t *testing.T) {
|
||||
aUp := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(500)
|
||||
fmt.Fprint(w, `{"error":"boom"}`)
|
||||
}))
|
||||
defer aUp.Close()
|
||||
bUp := mockUpstream()
|
||||
defer bUp.Close()
|
||||
g := newTestGateway(t,
|
||||
// Both sources expose the same model id so the direct scheduler has a
|
||||
// real fallback candidate after a hard-fails; b is the healthy one.
|
||||
config.Source{Name: "a", BaseURL: aUp.URL, Adapter: "openai", Models: []config.Model{{ID: "m"}}},
|
||||
config.Source{Name: "b", BaseURL: bUp.URL, Adapter: "openai", Models: []config.Model{{ID: "m"}}},
|
||||
)
|
||||
body := doReq(t, g, "POST", "/v1/chat/completions",
|
||||
`{"model":"m","stream":true,"messages":[{"role":"user","content":"hi"}]}`)
|
||||
if body.Code != 200 {
|
||||
t.Fatalf("status=%d body=%s", body.Code, body.Body.String())
|
||||
}
|
||||
sb := body.Body.String()
|
||||
if !strings.Contains(sb, `"finish_reason":"stop"`) || !strings.Contains(sb, "[DONE]") {
|
||||
t.Fatalf("stream incomplete: %s", sb)
|
||||
}
|
||||
recs := g.stats.AuditRecords(0, 0, "")
|
||||
if len(recs) == 0 {
|
||||
t.Fatal("no audit records")
|
||||
}
|
||||
last := recs[len(recs)-1]
|
||||
if last.Source != "b" || !last.OK {
|
||||
t.Fatalf("audit rec source=%q ok=%v, want source=b ok=true (actual serving source)", last.Source, last.OK)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user