mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-21 05:57:53 +00:00
fix(gateway): close billing, limit, routing and translation gaps
Billing: - Passthrough SSE streams now bill the usage the upstream reports. The gateway falls back to a request-size estimate only on a 2xx stream or a truncated JSON body. - A unified stream that fails after it starts is estimated, not billed at zero. - Billing uses the built-in LiteLLM price table and honours modelMeta free: true. - Pricing entries are validated. Cost arithmetic rounds the total and saturates on overflow. - Gemini thinking tokens count as completion tokens. Limits and storage: - A rejected request rolls back its counts on the other scopes, so one user cannot lock out a group. - Redis EXPIRE replies are checked. Keys recreated without a TTL are repaired at once. - A mutating pipeline is not retried after reply bytes arrive. - History, model totals and target usage report a store outage. Routing and translation: - Discovered ids with a provider-like prefix are listed only in owner/id form, so listing and routing agree. - Anthropic and Gemini discovery follows pagination. - The developer role maps to system. - stream_options.include_usage is honoured on translated streams. - tool_choice none keeps tools when history holds tool blocks. - refusal maps to content_filter. Data URI parameters are stripped. Streaming, retry and failover: - SSE accepts CR, LF and CRLF line ends. This closes a CR injection path. - Retry-After is never jittered. A cancelled wait returns the context error. - The watchdog timer stops at end of body. - Oversized bodies return 413. - A failed stream ends with an error event. - Failover restores pre-existing headers. Client cancels do not count as provider successes. Also: "admin" is a reserved provider name. metrics.path must start with "/". Pricing warnings go to the right gateway instance. The admin targets view sends access as null or [] so the two states differ. Upstream organisation headers are stripped. The users-file reload error is logged once per error. Dead code is removed, and test-only helpers move to test files. Claude-Session: https://claude.ai/code/session_01V5jYJNxpDFdJkyyExpADrv
This commit is contained in:
+43
-4
@@ -320,7 +320,10 @@ func geminiRequestFromOpenAI(req map[string]any) (map[string]any, error) {
|
||||
continue
|
||||
}
|
||||
switch role, _ := msg["role"].(string); role {
|
||||
case "system":
|
||||
case "system", "developer":
|
||||
// "developer" is the OpenAI o-series/gpt-5 replacement for
|
||||
// "system" — Gemini has no separate developer role, so it
|
||||
// folds into the same systemInstruction text.
|
||||
switch c := msg["content"].(type) {
|
||||
case string:
|
||||
systemParts = append(systemParts, c)
|
||||
@@ -451,9 +454,14 @@ type geminiCandidate struct {
|
||||
|
||||
// geminiUsageMetadata is the "usageMetadata" object on a Gemini
|
||||
// GenerateContentResponse, non-streaming or streaming alike.
|
||||
// ThoughtsTokenCount is Gemini's reasoning/thinking token spend (2.5+
|
||||
// thinking models) — it is billed as output but is not included in
|
||||
// CandidatesTokenCount, so safeCompletion folds it into the completion
|
||||
// total rather than leaving it undercounted.
|
||||
type geminiUsageMetadata struct {
|
||||
PromptTokenCount int64 `json:"promptTokenCount"`
|
||||
CandidatesTokenCount int64 `json:"candidatesTokenCount"`
|
||||
ThoughtsTokenCount int64 `json:"thoughtsTokenCount"`
|
||||
}
|
||||
|
||||
// geminiGenerateContentResponse is the subset of a Gemini
|
||||
@@ -627,12 +635,14 @@ func (m *geminiUsageMetadata) safePrompt() int64 {
|
||||
return m.PromptTokenCount
|
||||
}
|
||||
|
||||
// safeCompletion returns m.CandidatesTokenCount, or 0 when m is nil.
|
||||
// safeCompletion returns m.CandidatesTokenCount plus m.ThoughtsTokenCount
|
||||
// (thinking-model reasoning tokens, billed as output but reported
|
||||
// separately from CandidatesTokenCount), or 0 when m is nil.
|
||||
func (m *geminiUsageMetadata) safeCompletion() int64 {
|
||||
if m == nil {
|
||||
return 0
|
||||
}
|
||||
return m.CandidatesTokenCount
|
||||
return m.CandidatesTokenCount + m.ThoughtsTokenCount
|
||||
}
|
||||
|
||||
// geminiStreamState is one streaming chat completion's translation state:
|
||||
@@ -707,6 +717,35 @@ func (st *geminiStreamState) chunk(delta map[string]any, finishReason *string) [
|
||||
return b
|
||||
}
|
||||
|
||||
// usageChunk builds the final usage-only "chat.completion.chunk" payload —
|
||||
// empty "choices" plus a populated "usage" — mirroring the shape
|
||||
// provider_openai.go's forwardStream relays from a native OpenAI upstream
|
||||
// when stream_options.include_usage is set (M1 fix). Gemini's own stream
|
||||
// has no equivalent chunk, so the adapter calls this once, after the
|
||||
// stream ends, only when the client asked for usage; st.u already holds
|
||||
// the last-usageMetadata-wins totals by then.
|
||||
func (st *geminiStreamState) usageChunk() []byte {
|
||||
b, err := json.Marshal(map[string]any{
|
||||
"id": chatCompletionIDPrefix + st.id,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": st.created,
|
||||
"model": st.model,
|
||||
"choices": []any{},
|
||||
"usage": map[string]any{
|
||||
"prompt_tokens": st.u.prompt,
|
||||
"completion_tokens": st.u.completion,
|
||||
"total_tokens": st.u.total(),
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
// Built entirely from this file's own map[string]any/string/int64
|
||||
// literals plus st.u's int64 fields — never a value that can fail
|
||||
// to marshal.
|
||||
panic(fmt.Sprintf("llmgateway: gemini stream usage chunk failed to marshal: %v", err))
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
// translate maps one upstream Gemini SSE event to zero or more
|
||||
// OpenAI-shaped "chat.completion.chunk" payloads, per this file's stream
|
||||
// mapping table. Every event's data is a full GenerateContentResponse
|
||||
@@ -747,7 +786,7 @@ func (st *geminiStreamState) translate(ev sseEvent) ([][]byte, error) {
|
||||
}
|
||||
|
||||
if resp.UsageMetadata != nil {
|
||||
st.u = usage{prompt: resp.UsageMetadata.PromptTokenCount, completion: resp.UsageMetadata.CandidatesTokenCount}
|
||||
st.u = usage{prompt: resp.UsageMetadata.safePrompt(), completion: resp.UsageMetadata.safeCompletion()}
|
||||
}
|
||||
|
||||
var chunks [][]byte
|
||||
|
||||
Reference in New Issue
Block a user