Fix OpenAI compat usage parsing
This commit is contained in:
parent
6e66edbb09
commit
0393bd7c82
@ -306,6 +306,37 @@ type ResponsesUsage struct {
|
||||
OutputTokensDetails *ResponsesOutputTokensDetails `json:"output_tokens_details,omitempty"`
|
||||
}
|
||||
|
||||
func (u *ResponsesUsage) UnmarshalJSON(data []byte) error {
|
||||
type responsesUsageAlias ResponsesUsage
|
||||
var aux struct {
|
||||
responsesUsageAlias
|
||||
PromptTokens int `json:"prompt_tokens"`
|
||||
CompletionTokens int `json:"completion_tokens"`
|
||||
PromptTokensDetails *ResponsesInputTokensDetails `json:"prompt_tokens_details,omitempty"`
|
||||
CompletionTokensDetails *ResponsesOutputTokensDetails `json:"completion_tokens_details,omitempty"`
|
||||
}
|
||||
if err := json.Unmarshal(data, &aux); err != nil {
|
||||
return err
|
||||
}
|
||||
*u = ResponsesUsage(aux.responsesUsageAlias)
|
||||
if u.InputTokens == 0 && aux.PromptTokens != 0 {
|
||||
u.InputTokens = aux.PromptTokens
|
||||
}
|
||||
if u.OutputTokens == 0 && aux.CompletionTokens != 0 {
|
||||
u.OutputTokens = aux.CompletionTokens
|
||||
}
|
||||
if u.InputTokensDetails == nil && aux.PromptTokensDetails != nil {
|
||||
u.InputTokensDetails = aux.PromptTokensDetails
|
||||
}
|
||||
if u.OutputTokensDetails == nil && aux.CompletionTokensDetails != nil {
|
||||
u.OutputTokensDetails = aux.CompletionTokensDetails
|
||||
}
|
||||
if u.TotalTokens == 0 && (u.InputTokens != 0 || u.OutputTokens != 0) {
|
||||
u.TotalTokens = u.InputTokens + u.OutputTokens
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ResponsesInputTokensDetails breaks down input token usage.
|
||||
type ResponsesInputTokensDetails struct {
|
||||
CachedTokens int `json:"cached_tokens,omitempty"`
|
||||
|
||||
@ -183,6 +183,63 @@ func TestForwardAsAnthropic_NormalizesRoutingAndEffortForGpt54XHigh(t *testing.T
|
||||
t.Logf("response body: %s", rec.Body.String())
|
||||
}
|
||||
|
||||
func TestForwardAsAnthropic_MappedClaudeModelAcceptsChatUsageShape(t *testing.T) {
|
||||
t.Parallel()
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
body := []byte(`{"model":"claude-opus-4-7","max_tokens":16,"messages":[{"role":"user","content":"compact this"}],"stream":true}`)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(body))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
upstreamBody := strings.Join([]string{
|
||||
`data: {"type":"response.created","response":{"id":"resp_compact","model":"gpt-5.5","status":"in_progress","output":[]}}`,
|
||||
"",
|
||||
`data: {"type":"response.output_text.delta","delta":"ok"}`,
|
||||
"",
|
||||
`data: {"type":"response.completed","response":{"id":"resp_compact","object":"response","model":"gpt-5.5","status":"completed","output":[{"type":"message","id":"msg_1","role":"assistant","status":"completed","content":[{"type":"output_text","text":"ok"}]}],"usage":{"prompt_tokens":31,"completion_tokens":9,"total_tokens":40,"prompt_tokens_details":{"cached_tokens":11}}}}`,
|
||||
"",
|
||||
"data: [DONE]",
|
||||
"",
|
||||
}, "\n")
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"text/event-stream"}, "x-request-id": []string{"rid_compact_usage"}},
|
||||
Body: io.NopCloser(strings.NewReader(upstreamBody)),
|
||||
}}
|
||||
|
||||
svc := &OpenAIGatewayService{
|
||||
httpUpstream: upstream,
|
||||
cfg: &config.Config{Security: config.SecurityConfig{URLAllowlist: config.URLAllowlistConfig{Enabled: false}}},
|
||||
}
|
||||
account := &Account{
|
||||
ID: 1,
|
||||
Name: "openai-apikey",
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeAPIKey,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"api_key": "sk-test",
|
||||
"base_url": "https://api.openai.com/v1",
|
||||
"model_mapping": map[string]any{
|
||||
"gpt-5.5": "gpt-5.5",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
result, err := svc.ForwardAsAnthropic(context.Background(), c, account, body, "", "gpt-5.5")
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, "claude-opus-4-7", result.Model)
|
||||
require.Equal(t, "gpt-5.5", result.BillingModel)
|
||||
require.Equal(t, "gpt-5.5", result.UpstreamModel)
|
||||
require.Equal(t, 31, result.Usage.InputTokens)
|
||||
require.Equal(t, 9, result.Usage.OutputTokens)
|
||||
require.Equal(t, 11, result.Usage.CacheReadInputTokens)
|
||||
require.Equal(t, "gpt-5.5", gjson.GetBytes(upstream.lastBody, "model").String())
|
||||
}
|
||||
|
||||
func TestForwardAsAnthropic_InjectsPromptCacheKeyForAPIKeyMessagesDispatch(t *testing.T) {
|
||||
t.Parallel()
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
@ -4639,28 +4639,47 @@ func (s *OpenAIGatewayService) parseSSEUsageBytes(data []byte, usage *OpenAIUsag
|
||||
return
|
||||
}
|
||||
|
||||
usage.InputTokens = int(gjson.GetBytes(data, "response.usage.input_tokens").Int())
|
||||
usage.OutputTokens = int(gjson.GetBytes(data, "response.usage.output_tokens").Int())
|
||||
usage.CacheReadInputTokens = int(gjson.GetBytes(data, "response.usage.input_tokens_details.cached_tokens").Int())
|
||||
usage.ImageOutputTokens = int(gjson.GetBytes(data, "response.usage.output_tokens_details.image_tokens").Int())
|
||||
if parsedUsage, ok := extractOpenAIUsageFromJSONBytes(data); ok {
|
||||
*usage = parsedUsage
|
||||
}
|
||||
}
|
||||
|
||||
func extractOpenAIUsageFromJSONBytes(body []byte) (OpenAIUsage, bool) {
|
||||
if len(body) == 0 || !gjson.ValidBytes(body) {
|
||||
return OpenAIUsage{}, false
|
||||
}
|
||||
values := gjson.GetManyBytes(
|
||||
body,
|
||||
"usage.input_tokens",
|
||||
"usage.output_tokens",
|
||||
"usage.input_tokens_details.cached_tokens",
|
||||
"usage.output_tokens_details.image_tokens",
|
||||
)
|
||||
if usage, ok := openAIUsageFromGJSON(gjson.GetBytes(body, "usage")); ok {
|
||||
return usage, true
|
||||
}
|
||||
return openAIUsageFromGJSON(gjson.GetBytes(body, "response.usage"))
|
||||
}
|
||||
|
||||
func openAIUsageFromGJSON(value gjson.Result) (OpenAIUsage, bool) {
|
||||
if !value.Exists() || !value.IsObject() {
|
||||
return OpenAIUsage{}, false
|
||||
}
|
||||
inputTokens := value.Get("input_tokens").Int()
|
||||
if inputTokens == 0 {
|
||||
inputTokens = value.Get("prompt_tokens").Int()
|
||||
}
|
||||
outputTokens := value.Get("output_tokens").Int()
|
||||
if outputTokens == 0 {
|
||||
outputTokens = value.Get("completion_tokens").Int()
|
||||
}
|
||||
cacheReadTokens := value.Get("input_tokens_details.cached_tokens").Int()
|
||||
if cacheReadTokens == 0 {
|
||||
cacheReadTokens = value.Get("prompt_tokens_details.cached_tokens").Int()
|
||||
}
|
||||
imageOutputTokens := value.Get("output_tokens_details.image_tokens").Int()
|
||||
if imageOutputTokens == 0 {
|
||||
imageOutputTokens = value.Get("completion_tokens_details.image_tokens").Int()
|
||||
}
|
||||
return OpenAIUsage{
|
||||
InputTokens: int(values[0].Int()),
|
||||
OutputTokens: int(values[1].Int()),
|
||||
CacheReadInputTokens: int(values[2].Int()),
|
||||
ImageOutputTokens: int(values[3].Int()),
|
||||
InputTokens: int(inputTokens),
|
||||
OutputTokens: int(outputTokens),
|
||||
CacheCreationInputTokens: int(value.Get("cache_creation_input_tokens").Int()),
|
||||
CacheReadInputTokens: int(cacheReadTokens),
|
||||
ImageOutputTokens: int(imageOutputTokens),
|
||||
}, true
|
||||
}
|
||||
|
||||
|
||||
@ -2174,6 +2174,25 @@ func TestParseSSEUsage_SelectiveParsing(t *testing.T) {
|
||||
require.Equal(t, 13, usage.InputTokens)
|
||||
require.Equal(t, 15, usage.OutputTokens)
|
||||
require.Equal(t, 4, usage.CacheReadInputTokens)
|
||||
|
||||
svc.parseSSEUsage(`{"type":"response.completed","response":{"usage":{"prompt_tokens":21,"completion_tokens":8,"prompt_tokens_details":{"cached_tokens":6}}}}`, usage)
|
||||
require.Equal(t, 21, usage.InputTokens)
|
||||
require.Equal(t, 8, usage.OutputTokens)
|
||||
require.Equal(t, 6, usage.CacheReadInputTokens)
|
||||
}
|
||||
|
||||
func TestExtractOpenAIUsageFromJSONBytes_AcceptsResponseAndChatUsageShapes(t *testing.T) {
|
||||
usage, ok := extractOpenAIUsageFromJSONBytes([]byte(`{"id":"resp_1","usage":{"input_tokens":3,"output_tokens":5,"input_tokens_details":{"cached_tokens":2}}}`))
|
||||
require.True(t, ok)
|
||||
require.Equal(t, 3, usage.InputTokens)
|
||||
require.Equal(t, 5, usage.OutputTokens)
|
||||
require.Equal(t, 2, usage.CacheReadInputTokens)
|
||||
|
||||
usage, ok = extractOpenAIUsageFromJSONBytes([]byte(`{"type":"response.completed","response":{"usage":{"prompt_tokens":13,"completion_tokens":7,"prompt_tokens_details":{"cached_tokens":4}}}}`))
|
||||
require.True(t, ok)
|
||||
require.Equal(t, 13, usage.InputTokens)
|
||||
require.Equal(t, 7, usage.OutputTokens)
|
||||
require.Equal(t, 4, usage.CacheReadInputTokens)
|
||||
}
|
||||
|
||||
func TestExtractCodexFinalResponse_SampleReplay(t *testing.T) {
|
||||
|
||||
@ -399,15 +399,9 @@ func parseOpenAIWSResponseUsageFromCompletedEvent(message []byte, usage *OpenAIU
|
||||
if usage == nil || len(message) == 0 {
|
||||
return
|
||||
}
|
||||
values := gjson.GetManyBytes(
|
||||
message,
|
||||
"response.usage.input_tokens",
|
||||
"response.usage.output_tokens",
|
||||
"response.usage.input_tokens_details.cached_tokens",
|
||||
)
|
||||
usage.InputTokens = int(values[0].Int())
|
||||
usage.OutputTokens = int(values[1].Int())
|
||||
usage.CacheReadInputTokens = int(values[2].Int())
|
||||
if parsedUsage, ok := extractOpenAIUsageFromJSONBytes(message); ok {
|
||||
*usage = parsedUsage
|
||||
}
|
||||
}
|
||||
|
||||
func parseOpenAIWSErrorEventFields(message []byte) (code string, errType string, errMessage string) {
|
||||
|
||||
@ -29,6 +29,14 @@ func TestParseOpenAIWSResponseUsageFromCompletedEvent(t *testing.T) {
|
||||
require.Equal(t, 11, usage.InputTokens)
|
||||
require.Equal(t, 7, usage.OutputTokens)
|
||||
require.Equal(t, 3, usage.CacheReadInputTokens)
|
||||
|
||||
parseOpenAIWSResponseUsageFromCompletedEvent(
|
||||
[]byte(`{"type":"response.completed","response":{"usage":{"prompt_tokens":19,"completion_tokens":5,"prompt_tokens_details":{"cached_tokens":4}}}}`),
|
||||
usage,
|
||||
)
|
||||
require.Equal(t, 19, usage.InputTokens)
|
||||
require.Equal(t, 5, usage.OutputTokens)
|
||||
require.Equal(t, 4, usage.CacheReadInputTokens)
|
||||
}
|
||||
|
||||
func TestOpenAIWSErrorEventHelpers_ConsistentWithWrapper(t *testing.T) {
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user