🔖 chore: add gpt-4o tiktoken

common/token.go

@@ -9,12 +9,13 @@ import (
 	"one-api/common/image"
 	"one-api/types"
 
-	"github.com/pkoukk/tiktoken-go"
+	"github.com/MartialBE/tiktoken-go"
 )
 
 var tokenEncoderMap = map[string]*tiktoken.Tiktoken{}
 var gpt35TokenEncoder *tiktoken.Tiktoken
 var gpt4TokenEncoder *tiktoken.Tiktoken
+var gpt4oTokenEncoder *tiktoken.Tiktoken
 
 func InitTokenEncoders() {
 	SysLog("initializing token encoders")
@@ -29,6 +30,11 @@ func InitTokenEncoders() {
 		FatalLog(fmt.Sprintf("failed to get gpt-4 token encoder: %s", err.Error()))
 	}
 
+	gpt4oTokenEncoder, err = tiktoken.EncodingForModel("gpt-4o")
+	if err != nil {
+		FatalLog(fmt.Sprintf("failed to get gpt-4o token encoder: %s", err.Error()))
+	}
+
 	SysLog("token encoders initialized")
 }
 
@@ -40,6 +46,8 @@ func getTokenEncoder(model string) *tiktoken.Tiktoken {
 
 	if strings.HasPrefix(model, "gpt-3.5") {
 		tokenEncoder = gpt35TokenEncoder
+	} else if strings.HasPrefix(model, "gpt-4o") {
+		tokenEncoder = gpt4oTokenEncoder
 	} else if strings.HasPrefix(model, "gpt-4") {
 		tokenEncoder = gpt4TokenEncoder
 	} else {

go.mod

@@ -4,6 +4,7 @@ module one-api
 go 1.18
 
 require (
+	github.com/MartialBE/tiktoken-go v0.1.7
 	github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.1
 	github.com/aws/smithy-go v1.20.1
 	github.com/gin-contrib/cors v1.7.0
@@ -18,7 +19,6 @@ require (
 	github.com/gomarkdown/markdown v0.0.0-20240328165702-4d01890c35c0
 	github.com/google/uuid v1.6.0
 	github.com/gorilla/websocket v1.5.1
-	github.com/pkoukk/tiktoken-go v0.1.6
 	github.com/shopspring/decimal v1.4.0
 	github.com/spf13/viper v1.18.2
 	github.com/stretchr/testify v1.9.0

go.sum

@@ -1,5 +1,7 @@
 filippo.io/edwards25519 v1.1.0 h1:FNf4tywRC1HmFuKW5xopWpigGjJKiJSV0Cqo0cJWDaA=
 filippo.io/edwards25519 v1.1.0/go.mod h1:BxyFTGdWcka3PhytdK4V28tE5sGfRvvvRV7EaN4VDT4=
+github.com/MartialBE/tiktoken-go v0.1.7 h1:QwYqiONL1+kc/bYtNIbfQ//bcRqyw8ETZ2HKKZTQv8A=
+github.com/MartialBE/tiktoken-go v0.1.7/go.mod h1:BTUpeOlungENRv13+nWc4IwWbBJvvR/8gOUTiXsrEXg=
 github.com/PaulSonOfLars/gotgbot/v2 v2.0.0-rc.25 h1:VCZg3OsKY19PcXBRRYk2ExeZ3mC8Hm4LqcXcINuFyY4=
 github.com/PaulSonOfLars/gotgbot/v2 v2.0.0-rc.25/go.mod h1:kL1v4iIjlalwm3gCYGvF4NLa3hs+aKEfRkNJvj4aoDU=
 github.com/antonlindstrom/pgstore v0.0.0-20200229204646-b08ebf1105e0/go.mod h1:2Ti6VUHVxpC0VSmTZzEvpzysnaGAfGBOoMIz5ykPyyw=
@@ -166,8 +168,6 @@ github.com/pelletier/go-toml/v2 v2.1.1 h1:LWAJwfNvjQZCFIDKWYQaM62NcYeYViCmWIwmOS
 github.com/pelletier/go-toml/v2 v2.1.1/go.mod h1:tJU2Z3ZkXwnxa4DPO899bsyIoywizdUvyaeZurnPPDc=
 github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
-github.com/pkoukk/tiktoken-go v0.1.6 h1:JF0TlJzhTbrI30wCvFuiw6FzP2+/bR+FIxUdgEAcUsw=
-github.com/pkoukk/tiktoken-go v0.1.6/go.mod h1:9NiV+i9mJKGj1rYOT+njbv+ZwA/zJxYdewGl6qVatpg=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U=
 github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=