feat: support usage extra unmarshaler

Author: mrh997

common.go

@@ -1,14 +1,47 @@
 package openai
 
+import (
+	"encoding/json"
+	"fmt"
+	"reflect"
+
+	openai "github.com/meguminnnnnnnnn/go-openai/internal"
+)
+
 // common.go defines common types used throughout the OpenAI API.
 
 // Usage Represents the total token usage per request to OpenAI.
 type Usage struct {
-	PromptTokens            int                      `json:"prompt_tokens"`
-	CompletionTokens        int                      `json:"completion_tokens"`
-	TotalTokens             int                      `json:"total_tokens"`
-	PromptTokensDetails     *PromptTokensDetails     `json:"prompt_tokens_details"`
-	CompletionTokensDetails *CompletionTokensDetails `json:"completion_tokens_details"`
+	PromptTokens            int                        `json:"prompt_tokens"`
+	CompletionTokens        int                        `json:"completion_tokens"`
+	TotalTokens             int                        `json:"total_tokens"`
+	PromptTokensDetails     *PromptTokensDetails       `json:"prompt_tokens_details"`
+	CompletionTokensDetails *CompletionTokensDetails   `json:"completion_tokens_details"`
+	ExtraFields             map[string]json.RawMessage `json:"-"`
+}
+
+func (u *Usage) UnmarshalJSON(data []byte) error {
+	if u == nil {
+		return fmt.Errorf("usage is nil")
+	}
+
+	type Alias Usage
+	alias := &Alias{}
+	err := json.Unmarshal(data, alias)
+	if err != nil {
+		return err
+	}
+
+	*u = Usage(*alias)
+
+	extra, err := openai.UnmarshalExtraFields(reflect.TypeOf(u), data)
+	if err != nil {
+		return err
+	}
+
+	u.ExtraFields = extra
+
+	return nil
 }
 
 // CompletionTokensDetails Breakdown of tokens used in a completion.

common_test.go

@@ -0,0 +1,41 @@
+package openai_test
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/meguminnnnnnnnn/go-openai"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestUsageUnmarshalJSON(t *testing.T) {
+	data := []byte(`{
+		"prompt_tokens": 10,
+		"completion_tokens": 20,
+		"total_tokens": 30,
+		"prompt_tokens_details": {
+			"cached_tokens": 15
+		},
+		"completion_tokens_details": {
+			"audio_tokens": 10
+		},
+		"extra_field": "extra_value"
+	}`)
+
+	usage := &openai.Usage{}
+	err := json.Unmarshal(data, usage)
+	assert.NoError(t, err)
+	assert.Equal(t, 10, usage.PromptTokens)
+	assert.Equal(t, 20, usage.CompletionTokens)
+	assert.Equal(t, 30, usage.TotalTokens)
+	assert.NotNil(t, usage.PromptTokensDetails)
+	assert.Equal(t, 15, usage.PromptTokensDetails.CachedTokens)
+	assert.NotNil(t, usage.CompletionTokensDetails)
+	assert.Equal(t, 10, usage.CompletionTokensDetails.AudioTokens)
+	assert.Len(t, usage.ExtraFields, 1)
+
+	var extraValue string
+	err = json.Unmarshal(usage.ExtraFields["extra_field"], &extraValue)
+	assert.NoError(t, err)
+	assert.Equal(t, "extra_value", extraValue)
+}

go.mod

@@ -2,6 +2,20 @@ module github.com/meguminnnnnnnnn/go-openai
 
 go 1.18
 
-require github.com/evanphx/json-patch v0.5.2
+require (
+	github.com/bytedance/sonic v1.14.0
+	github.com/evanphx/json-patch v0.5.2
+	github.com/stretchr/testify v1.10.0
+)
 
-require github.com/pkg/errors v0.9.1 // indirect
+require (
+	github.com/bytedance/sonic/loader v0.3.0 // indirect
+	github.com/cloudwego/base64x v0.1.5 // indirect
+	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/klauspost/cpuid/v2 v2.0.9 // indirect
+	github.com/pkg/errors v0.9.1 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	github.com/twitchyliquid64/golang-asm v0.15.1 // indirect
+	golang.org/x/arch v0.0.0-20210923205945-b76863e36670 // indirect
+	gopkg.in/yaml.v3 v3.0.1 // indirect
+)

go.sum

@@ -1,5 +1,40 @@
+github.com/bytedance/sonic v1.14.0 h1:/OfKt8HFw0kh2rj8N0F6C/qPGRESq0BbaNZgcNXXzQQ=
+github.com/bytedance/sonic v1.14.0/go.mod h1:WoEbx8WTcFJfzCe0hbmyTGrfjt8PzNEBdxlNUO24NhA=
+github.com/bytedance/sonic/loader v0.1.1/go.mod h1:ncP89zfokxS5LZrJxl5z0UJcsk4M4yY2JpfqGeCtNLU=
+github.com/bytedance/sonic/loader v0.3.0 h1:dskwH8edlzNMctoruo8FPTJDF3vLtDT0sXZwvZJyqeA=
+github.com/bytedance/sonic/loader v0.3.0/go.mod h1:N8A3vUdtUebEY2/VQC0MyhYeKUFosQU6FxH2JmUe6VI=
+github.com/cloudwego/base64x v0.1.5 h1:XPciSp1xaq2VCSt6lF0phncD4koWyULpl5bUxbfCyP4=
+github.com/cloudwego/base64x v0.1.5/go.mod h1:0zlkT4Wn5C6NdauXdJRhSKRlJvmclQ1hhJgA0rcu/8w=
+github.com/cloudwego/iasm v0.2.0/go.mod h1:8rXZaNYT2n95jn+zTI1sDr+IgcD2GVs0nlbbQPiEFhY=
+github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/evanphx/json-patch v0.5.2 h1:xVCHIVMUu1wtM/VkR9jVZ45N3FhZfYMMYGorLCR8P3k=
 github.com/evanphx/json-patch v0.5.2/go.mod h1:ZWS5hhDbVDyob71nXKNL0+PWn6ToqBHMikGIFbs31qQ=
 github.com/jessevdk/go-flags v1.4.0/go.mod h1:4FA24M0QyGHXBuZZK/XkWh8h0e1EYbRYJSGM75WSRxI=
+github.com/klauspost/cpuid/v2 v2.0.9 h1:lgaqFMSdTdQYdZ04uHyN2d/eKdOMyi2YLSvlQIBFYa4=
+github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg=
+github.com/knz/go-libedit v1.10.1/go.mod h1:MZTVkCWyz0oBc7JOWP3wNAzd002ZbM/5hgShxwh4x8M=
 github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
+github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
+github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
+github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
+github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA=
+github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
+github.com/twitchyliquid64/golang-asm v0.15.1 h1:SU5vSMR7hnwNxj24w34ZyCi/FmDZTkS4MhqMhdFk5YI=
+github.com/twitchyliquid64/golang-asm v0.15.1/go.mod h1:a1lVb/DtPvCB8fslRZhAngC2+aY1QWCk3Cedj/Gdt08=
+golang.org/x/arch v0.0.0-20210923205945-b76863e36670 h1:18EFjUmQOcUvxNYSkA6jO9VAiXCnxFY6NyDX0bHDmkU=
+golang.org/x/arch v0.0.0-20210923205945-b76863e36670/go.mod h1:5om86z9Hs0C8fWVUuoMHwpExlXzs5Tkyp9hOrfG7pp8=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
+gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
+nullprogram.com/x/optparse v1.0.0/go.mod h1:KdyPE+Igbe0jQUrVfMqDMeJQIJZEuyV7pjYmp6pbG50=

internal/unmarshaler.go

@@ -2,6 +2,10 @@ package openai
 
 import (
 	"encoding/json"
+	"fmt"
+	"reflect"
+
+	"github.com/bytedance/sonic"
 )
 
 type Unmarshaler interface {
@@ -13,3 +17,39 @@ type JSONUnmarshaler struct{}
 func (jm *JSONUnmarshaler) Unmarshal(data []byte, v any) error {
 	return json.Unmarshal(data, v)
 }
+
+func UnmarshalExtraFields(typ reflect.Type, data []byte) (map[string]json.RawMessage, error) {
+	m := make(map[string]json.RawMessage)
+	if err := sonic.Unmarshal(data, &m); err != nil {
+		return nil, err
+	}
+
+	for typ.Kind() == reflect.Ptr {
+		typ = typ.Elem()
+	}
+
+	if typ.Kind() != reflect.Struct {
+		return nil, fmt.Errorf("type is not a struct")
+	}
+
+	for i := 0; i < typ.NumField(); i++ {
+		field := typ.Field(i)
+
+		jsonTag := field.Tag.Get("json")
+		if jsonTag != "" {
+			delete(m, jsonTag)
+		} else {
+			if !field.IsExported() {
+				continue
+			}
+			delete(m, field.Name)
+		}
+	}
+
+	extra := make(map[string]json.RawMessage, len(m))
+	for k, v := range m {
+		extra[k] = v
+	}
+
+	return extra, nil
+}

internal/unmarshaler_test.go

@@ -1,10 +1,13 @@
 package openai_test
 
 import (
+	"encoding/json"
+	"reflect"
 	"testing"
 
 	openai "github.com/meguminnnnnnnnn/go-openai/internal"
 	"github.com/meguminnnnnnnnn/go-openai/internal/test/checks"
+	"github.com/stretchr/testify/assert"
 )
 
 func TestJSONUnmarshaler_Normal(t *testing.T) {
@@ -35,3 +38,25 @@ func TestJSONUnmarshaler_EmptyInput(t *testing.T) {
 	err := jm.Unmarshal(nil, &v)
 	checks.HasError(t, err, "should return error for nil input")
 }
+
+func TestUnmarshalExtraFields(t *testing.T) {
+	type TestStruct struct {
+		Field1 string `json:"field1"`
+		Field2 int
+		Field3 struct {
+			Field4 string `json:"field4"`
+		} `json:"field3"`
+	}
+
+	testData := []byte(`{"field1":"value1","Field2":2,"field3":{"field4":"value4"},"extraField1":"extraValue1"}`)
+	testStruct := &TestStruct{}
+	extra, err := openai.UnmarshalExtraFields(reflect.TypeOf(testStruct), testData)
+	assert.NoError(t, err)
+	assert.Len(t, extra, 1)
+
+	var extraValue1 string
+	err = json.Unmarshal(extra["extraField1"], &extraValue1)
+	assert.NoError(t, err)
+
+	assert.Equal(t, "extraValue1", extraValue1)
+}