diff --git a/go.mod b/go.mod index 69ac5273..cb5bd6ac 100644 --- a/go.mod +++ b/go.mod @@ -13,9 +13,10 @@ require ( github.com/joho/godotenv v1.3.0 github.com/json-iterator/go v1.1.12 github.com/mojocn/base64Captcha v1.3.5 + github.com/pkoukk/tiktoken-go v0.1.0 github.com/rhysd/go-github-selfupdate v1.2.3 github.com/spf13/cobra v1.2.1 - github.com/stretchr/testify v1.7.1 + github.com/stretchr/testify v1.8.2 github.com/xuri/excelize/v2 v2.5.0 github.com/yaoapp/gou v0.10.3 github.com/yaoapp/kun v0.9.0 @@ -24,6 +25,7 @@ require ( ) require ( + github.com/dlclark/regexp2 v1.8.1 // indirect github.com/google/go-github/v30 v30.1.0 // indirect github.com/google/go-querystring v1.0.0 // indirect github.com/inconshreveable/go-update v0.0.0-20160112193335-8152e7eb6ccf // indirect diff --git a/go.sum b/go.sum index 39142551..78a988f2 100644 --- a/go.sum +++ b/go.sum @@ -76,6 +76,8 @@ 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/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78= github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc= +github.com/dlclark/regexp2 v1.8.1 h1:6Lcdwya6GjPUNsBct8Lg/yRPwMhABj269AAzdGSiR+0= +github.com/dlclark/regexp2 v1.8.1/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8= github.com/elazarl/go-bindata-assetfs v1.0.1 h1:m0kkaHRKEu7tUIUFVwhGGGYClXvyl4RE03qmvRTNfbw= github.com/elazarl/go-bindata-assetfs v1.0.1/go.mod h1:v+YaWX3bdea5J/mo8dSETolEo7R71Vk1u8bnjau5yw4= github.com/envoyproxy/go-control-plane v0.9.0/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4= @@ -320,6 +322,8 @@ github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINE github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pkg/sftp v1.10.1/go.mod h1:lYOWFsE0bwd1+KfKJaKeuokY15vzFx25BLbzYYoAxZI= +github.com/pkoukk/tiktoken-go v0.1.0 h1:X1uP3+Nd8C3xe6AIGRWjchrylyaye0FDDTG22cxNQZs= +github.com/pkoukk/tiktoken-go v0.1.0/go.mod h1:BijIqAP84FMYC4XbdJgjyMpiSjusU8x0Y0W9K2t0QtU= 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/posener/complete v1.1.1/go.mod h1:em0nMJCgc9GFtwrmVmEMR/ZL6WyhyjMBndrE9hABlRI= @@ -351,14 +355,18 @@ github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA= github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= github.com/spf13/viper v1.8.1/go.mod h1:o0Pch8wJ9BVSWGQMbra6iw0oQ5oktSIBaujf1rJH9Ns= 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.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs= github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4= github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5cxcmMvtA= github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= -github.com/stretchr/testify v1.7.1 h1:5TQK59W5E3v0r2duFAb7P95B6hEeOyEnHRa8MjYSMTY= 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.2 h1:+h33VjcLVPDHtOdpUCuF+7gSuG3yGIftsP1YvFihtJ8= +github.com/stretchr/testify v1.8.2/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4= github.com/subosito/gotenv v1.2.0/go.mod h1:N0PQaV/YGNqwC0u51sEeR/aUtSLEXKX9iv69rRypqCw= github.com/tcnksm/go-gitconfig v0.1.2 h1:iiDhRitByXAEyjgBqsKi9QU4o2TNtv9kPP3RgPgXBPw= github.com/tcnksm/go-gitconfig v0.1.2/go.mod h1:/8EhP4H7oJZdIPyT+/UIsG87kTzrzM4UsLGSItWYCpE= diff --git a/main.go b/main.go index 7ded00d2..b9bed45a 100644 --- a/main.go +++ b/main.go @@ -6,6 +6,7 @@ import ( _ "github.com/yaoapp/gou/encoding" _ "github.com/yaoapp/yao/crypto" _ "github.com/yaoapp/yao/helper" + _ "github.com/yaoapp/yao/openai" _ "github.com/yaoapp/yao/utils" ) diff --git a/openai/openai.go b/openai/openai.go new file mode 100644 index 00000000..22ec932a --- /dev/null +++ b/openai/openai.go @@ -0,0 +1,13 @@ +package openai + +import "github.com/pkoukk/tiktoken-go" + +// Tiktoken get number of tokens +func Tiktoken(model string, input string) (int, error) { + tkm, err := tiktoken.EncodingForModel(model) + if err != nil { + return 0, err + } + token := tkm.Encode(input, nil, nil) + return len(token), nil +} diff --git a/openai/process.go b/openai/process.go new file mode 100644 index 00000000..bad5f911 --- /dev/null +++ b/openai/process.go @@ -0,0 +1,22 @@ +package openai + +import ( + "github.com/yaoapp/gou/process" + "github.com/yaoapp/kun/exception" +) + +func init() { + process.Register("yao.openai.tiktoken", ProcessTiktoken) +} + +// ProcessTiktoken get number of tokens +func ProcessTiktoken(process *process.Process) interface{} { + process.ValidateArgNums(2) + model := process.ArgsString(0) + input := process.ArgsString(1) + nums, err := Tiktoken(model, input) + if err != nil { + exception.New("Tiktoken error: %s", 400, err).Throw() + } + return nums +} diff --git a/openai/process_test.go b/openai/process_test.go new file mode 100644 index 00000000..d653f0fe --- /dev/null +++ b/openai/process_test.go @@ -0,0 +1,19 @@ +package openai + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/yaoapp/gou/process" +) + +func TestTiktoken(t *testing.T) { + // Hash + args := []interface{}{"gpt-3.5-turbo", "hello world"} + res := process.New("yao.openai.Tiktoken", args...).Run() + assert.Equal(t, 2, res) + + args = []interface{}{"gpt-3.5-turbo", "你好世界!"} + res = process.New("yao.openai.Tiktoken", args...).Run() + assert.Equal(t, 6, res) +}