BCAI/easyai-ai-gateway
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: add runner failover rules and cache affinity

Browse Source
This commit is contained in:
wangbo
2026-06-28 20:50:23 +08:00
parent 41bb3a6525
commit 0e675e259c
31 changed files with 2939 additions and 470 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/internal/clients/clients_test.go
+32 -1
View File
@@ -302,7 +302,7 @@ func TestOpenAIClientChatContract(t *testing.T) {
"choices": []any{map[string]any{
"message": map[string]any{"role": "assistant", "content": "ok"},
}},
"usage": map[string]any{"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5},
"usage": map[string]any{"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5, "prompt_cache_hit_tokens": 2},
})
}))
defer server.Close()
@@ -327,6 +327,14 @@ func TestOpenAIClientChatContract(t *testing.T) {
if response.Usage.TotalTokens != 5 || response.Result["id"] != "chatcmpl-test" {
t.Fatalf("unexpected response: %+v", response)
}
if response.Usage.CachedInputTokens != 2 {
t.Fatalf("expected cached input token usage, got %+v", response.Usage)
}
resultUsage, _ := response.Result["usage"].(map[string]any)
promptDetails, _ := resultUsage["prompt_tokens_details"].(map[string]any)
if intFromAny(promptDetails["cached_tokens"]) != 2 {
t.Fatalf("expected normalized cached prompt tokens in result usage, got %+v", resultUsage)
}
if response.RequestID != "req-chat-test" || response.ResponseStartedAt.IsZero() || response.ResponseFinishedAt.IsZero() {
t.Fatalf("response metadata was not captured: %+v", response)
}
@@ -335,6 +343,29 @@ func TestOpenAIClientChatContract(t *testing.T) {
}
}
func TestUsageFromOpenAIUsageTracksKnownCachedInputZero(t *testing.T) {
knownZero := usageFromOpenAIUsage(map[string]any{
"prompt_tokens": 100,
"completion_tokens": 10,
"total_tokens": 110,
"prompt_tokens_details": map[string]any{
"cached_tokens": 0,
},
})
if !knownZero.CachedInputTokensKnown || knownZero.CachedInputTokens != 0 {
t.Fatalf("expected known zero cached input tokens, got %+v", knownZero)
}
absent := usageFromOpenAIUsage(map[string]any{
"prompt_tokens": 100,
"completion_tokens": 10,
"total_tokens": 110,
})
if absent.CachedInputTokensKnown {
t.Fatalf("expected missing cached field to remain unknown, got %+v", absent)
}
}
func TestOpenAIClientEmbeddingsContract(t *testing.T) {
var gotPath string
var gotModel string
apps/api/internal/clients/gemini.go
+2 -2
View File
@@ -680,11 +680,11 @@ func geminiUsage(raw map[string]any) Usage {
input := intFromAny(usageMap["prompt_tokens"])
output := intFromAny(usageMap["completion_tokens"])
total := intFromAny(usageMap["total_tokens"])
cachedInput := cachedInputTokensFromOpenAIUsage(usageMap)
cachedInput, cachedInputKnown := cachedInputTokensValueFromOpenAIUsage(usageMap)
if cachedInput > input && input > 0 {
cachedInput = input
}
return Usage{InputTokens: input, OutputTokens: output, CachedInputTokens: cachedInput, TotalTokens: total}
return Usage{InputTokens: input, OutputTokens: output, CachedInputTokens: cachedInput, CachedInputTokensKnown: cachedInputKnown, TotalTokens: total}
}
func geminiUsageMap(raw map[string]any) map[string]any {
apps/api/internal/clients/helpers.go
+58 -7
View File
@@ -259,6 +259,9 @@ func NormalizeChatCompletionResult(result map[string]any) map[string]any {
if result == nil {
return nil
}
if usage, ok := result["usage"].(map[string]any); ok && len(usage) > 0 {
result["usage"] = NormalizeChatCompletionUsage(usage)
}
choices, _ := result["choices"].([]any)
for _, rawChoice := range choices {
choice, _ := rawChoice.(map[string]any)
@@ -280,6 +283,9 @@ func NormalizeChatCompletionStreamEvent(event map[string]any) map[string]any {
if event == nil {
return nil
}
if usage, ok := event["usage"].(map[string]any); ok && len(usage) > 0 {
event["usage"] = NormalizeChatCompletionUsage(usage)
}
choices, _ := event["choices"].([]any)
for _, rawChoice := range choices {
choice, _ := rawChoice.(map[string]any)
@@ -295,6 +301,33 @@ func NormalizeChatCompletionStreamEvent(event map[string]any) map[string]any {
return event
}
func NormalizeChatCompletionUsage(usage map[string]any) map[string]any {
if len(usage) == 0 {
return usage
}
normalized := usageFromOpenAIUsage(usage)
out := cloneMapAny(usage)
if normalized.InputTokens > 0 && intFromAny(out["prompt_tokens"]) == 0 {
out["prompt_tokens"] = normalized.InputTokens
}
if normalized.OutputTokens > 0 && intFromAny(out["completion_tokens"]) == 0 {
out["completion_tokens"] = normalized.OutputTokens
}
if normalized.TotalTokens > 0 && intFromAny(out["total_tokens"]) == 0 {
out["total_tokens"] = normalized.TotalTokens
}
if normalized.CachedInputTokens > 0 {
promptDetails, _ := firstPresent(out["prompt_tokens_details"], out["promptTokensDetails"]).(map[string]any)
normalizedPromptDetails := cloneMapAny(promptDetails)
if normalizedPromptDetails == nil {
normalizedPromptDetails = map[string]any{}
}
normalizedPromptDetails["cached_tokens"] = normalized.CachedInputTokens
out["prompt_tokens_details"] = normalizedPromptDetails
}
return out
}
func normalizeToolCallsContainer(container map[string]any, stream bool) {
if container == nil {
return
@@ -877,27 +910,36 @@ func sortedStreamToolCalls(toolCalls map[int]map[string]any) []any {
func usageFromOpenAI(result map[string]any) Usage {
usage, _ := result["usage"].(map[string]any)
return usageFromOpenAIUsage(usage)
}
func usageFromOpenAIUsage(usage map[string]any) Usage {
input := intFromAny(firstPresent(usage["prompt_tokens"], usage["input_tokens"]))
output := intFromAny(firstPresent(usage["completion_tokens"], usage["output_tokens"]))
total := intFromAny(usage["total_tokens"])
if total == 0 {
total = input + output
}
cachedInput := cachedInputTokensFromOpenAIUsage(usage)
cachedInput, cachedInputKnown := cachedInputTokensValueFromOpenAIUsage(usage)
if cachedInput > input && input > 0 {
cachedInput = input
}
return Usage{InputTokens: input, OutputTokens: output, CachedInputTokens: cachedInput, TotalTokens: total}
return Usage{InputTokens: input, OutputTokens: output, CachedInputTokens: cachedInput, CachedInputTokensKnown: cachedInputKnown, TotalTokens: total}
}
func cachedInputTokensFromOpenAIUsage(usage map[string]any) int {
cachedInput, _ := cachedInputTokensValueFromOpenAIUsage(usage)
return cachedInput
}
func cachedInputTokensValueFromOpenAIUsage(usage map[string]any) (int, bool) {
if len(usage) == 0 {
return 0
return 0, false
}
promptDetails, _ := firstPresent(usage["prompt_tokens_details"], usage["promptTokensDetails"]).(map[string]any)
inputDetails, _ := firstPresent(usage["input_tokens_details"], usage["inputTokensDetails"]).(map[string]any)
usageMetadata, _ := usage["usageMetadata"].(map[string]any)
return intFromAny(firstPresent(
value, ok := firstPresentValue(
promptDetails["cached_tokens"],
promptDetails["cachedTokens"],
promptDetails["cache_read_input_tokens"],
@@ -918,7 +960,11 @@ func cachedInputTokensFromOpenAIUsage(usage map[string]any) int {
usage["cachedContentTokenCount"],
usageMetadata["cached_content_token_count"],
usageMetadata["cachedContentTokenCount"],
))
)
if !ok {
return 0, false
}
return intFromAny(value), true
}
func requestIDFromHTTPResponse(resp *http.Response) string {
@@ -979,12 +1025,17 @@ func firstNonEmptyString(values ...any) string {
}
func firstPresent(values ...any) any {
value, _ := firstPresentValue(values...)
return value
}
func firstPresentValue(values ...any) (any, bool) {
for _, value := range values {
if value != nil {
return value
return value, true
}
}
return nil
return nil, false
}
func errorMessage(raw []byte, fallback string) string {
apps/api/internal/clients/simulation.go
+56 -2
View File
@@ -107,6 +107,7 @@ func simulationProfile(request Request) string {
}
func simulatedResult(request Request) map[string]any {
usage := simulatedUsage(request)
switch request.Kind {
case "chat.completions":
return map[string]any{
@@ -122,7 +123,7 @@ func simulatedResult(request Request) map[string]any {
"content": fmt.Sprintf("simulation response from %s", request.Candidate.Provider),
},
}},
"usage": map[string]any{"prompt_tokens": 12, "completion_tokens": 8, "total_tokens": 20},
"usage": simulatedOpenAIUsageMap(usage),
}
case "responses":
return map[string]any{
@@ -131,7 +132,7 @@ func simulatedResult(request Request) map[string]any {
"created_at": nowUnix(),
"model": request.Model,
"output_text": fmt.Sprintf("simulation response from %s", request.Candidate.Provider),
"usage": map[string]any{"input_tokens": 12, "output_tokens": 8, "total_tokens": 20},
"usage": simulatedResponseUsageMap(usage),
}
case "embeddings":
return simulatedEmbeddingResult(request)
@@ -373,6 +374,9 @@ func simulatedAudioData(request Request, fallbackPrompt string) []any {
}
func simulatedUsage(request Request) Usage {
if usage, ok := simulationUsageOverride(request); ok {
return usage
}
if request.ModelType == "chat" || request.ModelType == "text_generate" || request.Kind == "responses" {
return Usage{InputTokens: 12, OutputTokens: 8, TotalTokens: 20}
}
@@ -382,6 +386,56 @@ func simulatedUsage(request Request) Usage {
return Usage{}
}
func simulationUsageOverride(request Request) (Usage, bool) {
raw, _ := request.Body["simulationUsage"].(map[string]any)
if len(raw) == 0 {
raw, _ = request.Body["testUsage"].(map[string]any)
}
if len(raw) == 0 {
return Usage{}, false
}
usage := Usage{
InputTokens: intFromAny(firstPresent(raw["inputTokens"], raw["input_tokens"], raw["promptTokens"], raw["prompt_tokens"])),
OutputTokens: intFromAny(firstPresent(raw["outputTokens"], raw["output_tokens"], raw["completionTokens"], raw["completion_tokens"])),
CachedInputTokens: intFromAny(firstPresent(raw["cachedInputTokens"], raw["cached_input_tokens"], raw["cachedPromptTokens"], raw["cached_tokens"])),
TotalTokens: intFromAny(firstPresent(raw["totalTokens"], raw["total_tokens"])),
}
if _, ok := firstPresentValue(raw["cachedInputTokens"], raw["cached_input_tokens"], raw["cachedPromptTokens"], raw["cached_tokens"]); ok {
usage.CachedInputTokensKnown = true
}
if usage.TotalTokens == 0 {
usage.TotalTokens = usage.InputTokens + usage.OutputTokens
}
if usage.CachedInputTokens > usage.InputTokens && usage.InputTokens > 0 {
usage.CachedInputTokens = usage.InputTokens
}
return usage, true
}
func simulatedOpenAIUsageMap(usage Usage) map[string]any {
out := map[string]any{
"prompt_tokens": usage.InputTokens,
"completion_tokens": usage.OutputTokens,
"total_tokens": usage.TotalTokens,
}
if usage.CachedInputTokensKnown {
out["prompt_tokens_details"] = map[string]any{"cached_tokens": usage.CachedInputTokens}
}
return out
}
func simulatedResponseUsageMap(usage Usage) map[string]any {
out := map[string]any{
"input_tokens": usage.InputTokens,
"output_tokens": usage.OutputTokens,
"total_tokens": usage.TotalTokens,
}
if usage.CachedInputTokensKnown {
out["input_tokens_details"] = map[string]any{"cached_tokens": usage.CachedInputTokens}
}
return out
}
func simulatedProgress(request Request) []Progress {
provider := request.Candidate.Provider
if provider == "" {
apps/api/internal/clients/types.go
+5 -4
View File
@@ -36,10 +36,11 @@ type Response struct {
}
type Usage struct {
InputTokens int
OutputTokens int
CachedInputTokens int
TotalTokens int
InputTokens int
OutputTokens int
CachedInputTokens int
CachedInputTokensKnown bool
TotalTokens int
}
type Progress struct {