mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-09-20 00:48:00 +00:00
fix: auto-generate prompt_tokens_details from prompt_cache_hit_tokens in MarshalJSON
When only the legacy DeepSeek fields (prompt_cache_hit_tokens) are set but the OpenAI-standard prompt_tokens_details is nil, MarshalJSON now auto- generates the nested object. This ensures dsh (which reads the standard format first) sees cache hit data regardless of which adapter format the upstream uses.
This commit is contained in:
@ -107,6 +107,14 @@ type PromptTokensDetails struct {
|
|||||||
// standard keys (prompt_tokens/completion_tokens/total_tokens). Standard
|
// standard keys (prompt_tokens/completion_tokens/total_tokens). Standard
|
||||||
// clients such as DSH and DevEco Code read the *_tokens fields.
|
// clients such as DSH and DevEco Code read the *_tokens fields.
|
||||||
func (t TokenUsage) MarshalJSON() ([]byte, error) {
|
func (t TokenUsage) MarshalJSON() ([]byte, error) {
|
||||||
|
// Auto-generate prompt_tokens_details from legacy DeepSeek fields when
|
||||||
|
// the upstream adapter only set the standalone hit count (openai.lua
|
||||||
|
// does this in Lua, but other adapters or the standardSSEChunk fallback
|
||||||
|
// may not). dsh reads prompt_tokens_details.cached_tokens first.
|
||||||
|
pdetails := t.PromptTokensDetails
|
||||||
|
if pdetails == nil && t.PromptCacheHit > 0 {
|
||||||
|
pdetails = &PromptTokensDetails{CachedTokens: t.PromptCacheHit}
|
||||||
|
}
|
||||||
return json.Marshal(struct {
|
return json.Marshal(struct {
|
||||||
PromptTokens int `json:"prompt_tokens"`
|
PromptTokens int `json:"prompt_tokens"`
|
||||||
CompletionTokens int `json:"completion_tokens"`
|
CompletionTokens int `json:"completion_tokens"`
|
||||||
@ -124,7 +132,7 @@ func (t TokenUsage) MarshalJSON() ([]byte, error) {
|
|||||||
Prompt: t.Prompt,
|
Prompt: t.Prompt,
|
||||||
Completion: t.Completion,
|
Completion: t.Completion,
|
||||||
Total: t.Total,
|
Total: t.Total,
|
||||||
PromptTokensDetails: t.PromptTokensDetails,
|
PromptTokensDetails: pdetails,
|
||||||
PromptCacheHitTokens: t.PromptCacheHit,
|
PromptCacheHitTokens: t.PromptCacheHit,
|
||||||
PromptCacheMissTokens: t.PromptCacheMiss,
|
PromptCacheMissTokens: t.PromptCacheMiss,
|
||||||
})
|
})
|
||||||
|
|||||||
70
internal/types/types_test.go
Normal file
70
internal/types/types_test.go
Normal file
@ -0,0 +1,70 @@
|
|||||||
|
package types
|
||||||
|
|
||||||
|
import (
|
||||||
|
"encoding/json"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestTokenUsageMarshalCacheFields(t *testing.T) {
|
||||||
|
// DeepSeek-style: prompt_cache_hit_tokens + prompt_tokens_details auto-generated
|
||||||
|
u := TokenUsage{
|
||||||
|
Prompt: 100,
|
||||||
|
Completion: 50,
|
||||||
|
Total: 150,
|
||||||
|
PromptCacheHit: 80,
|
||||||
|
PromptCacheMiss: 20,
|
||||||
|
}
|
||||||
|
b, _ := json.Marshal(u)
|
||||||
|
var out map[string]interface{}
|
||||||
|
json.Unmarshal(b, &out)
|
||||||
|
check := func(name string, want interface{}) {
|
||||||
|
got := out[name]
|
||||||
|
if got != want {
|
||||||
|
t.Errorf("%s = %v (want %v)", name, got, want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
check("prompt_tokens", 100.0)
|
||||||
|
check("completion_tokens", 50.0)
|
||||||
|
check("prompt_cache_hit_tokens", 80.0)
|
||||||
|
check("prompt_cache_miss_tokens", 20.0)
|
||||||
|
details := out["prompt_tokens_details"]
|
||||||
|
if details == nil {
|
||||||
|
t.Error("prompt_tokens_details is missing (should be set from prompt_cache_hit_tokens?)")
|
||||||
|
} else {
|
||||||
|
d := details.(map[string]interface{})
|
||||||
|
if d["cached_tokens"] != 80.0 {
|
||||||
|
t.Errorf("cached_tokens = %v (want 80)", d["cached_tokens"])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// OpenAI-style: prompt_tokens_details.cached_tokens
|
||||||
|
u2 := TokenUsage{
|
||||||
|
Prompt: 200,
|
||||||
|
Completion: 100,
|
||||||
|
Total: 300,
|
||||||
|
PromptTokensDetails: &PromptTokensDetails{
|
||||||
|
CachedTokens: 150,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
b2, _ := json.Marshal(u2)
|
||||||
|
json.Unmarshal(b2, &out)
|
||||||
|
check2 := func(name string, want interface{}) {
|
||||||
|
got := out[name]
|
||||||
|
if got != want {
|
||||||
|
t.Errorf("OpenAI %s = %v (want %v)", name, got, want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
check2("prompt_tokens", 200.0)
|
||||||
|
details2 := out["prompt_tokens_details"]
|
||||||
|
if details2 == nil {
|
||||||
|
t.Error("OpenAI prompt_tokens_details is missing")
|
||||||
|
} else {
|
||||||
|
d := details2.(map[string]interface{})
|
||||||
|
if d["cached_tokens"] != 150.0 {
|
||||||
|
t.Errorf("OpenAI cached_tokens = %v (want 150)", d["cached_tokens"])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
t.Logf("DeepSeek JSON: %s", string(b))
|
||||||
|
t.Logf("OpenAI JSON: %s", string(b2))
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user