11package openai
22
33import (
4+ "encoding/json"
5+
46 openai "github.com/gptscript-ai/chat-completion-client"
7+ "github.com/gptscript-ai/gptscript/pkg/types"
8+ "github.com/pkoukk/tiktoken-go"
9+ tiktoken_loader "github.com/pkoukk/tiktoken-go-loader"
510)
611
12+ func init () {
13+ tiktoken .SetBpeLoader (tiktoken_loader .NewOfflineLoader ())
14+ }
15+
716const DefaultMaxTokens = 128_000
817
918func decreaseTenPercent (maxTokens int ) int {
@@ -12,22 +21,26 @@ func decreaseTenPercent(maxTokens int) int {
1221}
1322
1423func getBudget (maxTokens int ) int {
15- if maxTokens = =0 {
24+ if maxTokens < =0 {
1625 return DefaultMaxTokens
1726 }
1827 return maxTokens
1928}
2029
21- func dropMessagesOverCount (maxTokens int , msgs []openai.ChatCompletionMessage ) (result []openai.ChatCompletionMessage ) {
30+ func dropMessagesOverCount (maxTokens , toolTokenCount int , msgs []openai.ChatCompletionMessage ) (result []openai.ChatCompletionMessage , err error ) {
2231 var (
2332 lastSystem int
2433 withinBudget int
25- budget = getBudget (maxTokens )
34+ budget = getBudget (maxTokens ) - toolTokenCount
2635 )
2736
2837 for i , msg := range msgs {
2938 if msg .Role == openai .ChatMessageRoleSystem {
30- budget -= countMessage (msg )
39+ count , err := countMessage (msg )
40+ if err != nil {
41+ return nil , err
42+ }
43+ budget -= count
3144 lastSystem = i
3245 result = append (result , msg )
3346 } else {
@@ -37,7 +50,11 @@ func dropMessagesOverCount(maxTokens int, msgs []openai.ChatCompletionMessage) (
3750
3851 for i := len (msgs ) - 1 ; i > lastSystem ; i -- {
3952 withinBudget = i
40- budget -= countMessage (msgs [i ])
53+ count , err := countMessage (msgs [i ])
54+ if err != nil {
55+ return nil , err
56+ }
57+ budget -= count
4158 if budget <= 0 {
4259 break
4360 }
@@ -54,22 +71,42 @@ func dropMessagesOverCount(maxTokens int, msgs []openai.ChatCompletionMessage) (
5471 if withinBudget == len (msgs )- 1 {
5572 // We are going to drop all non system messages, which seems useless, so just return them
5673 // all and let it fail
57- return msgs
74+ return msgs , nil
5875 }
5976
60- return append (result , msgs [withinBudget :]... )
77+ return append (result , msgs [withinBudget :]... ), nil
6178}
6279
63- func countMessage (msg openai.ChatCompletionMessage ) (count int ) {
64- count += len (msg .Role )
65- count += len (msg .Content )
80+ func countMessage (msg openai.ChatCompletionMessage ) (int , error ) {
81+ encoding , err := tiktoken .GetEncoding ("o200k_base" )
82+ if err != nil {
83+ return 0 , err
84+ }
85+
86+ count := len (encoding .Encode (msg .Role , nil , nil ))
87+ count += len (encoding .Encode (msg .Content , nil , nil ))
6688 for _ , content := range msg .MultiContent {
67- count += len (content .Text )
89+ count += len (encoding . Encode ( content .Text , nil , nil ) )
6890 }
6991 for _ , tool := range msg .ToolCalls {
70- count += len (tool .Function .Name )
71- count += len (tool .Function .Arguments )
92+ count += len (encoding . Encode ( tool .Function .Name , nil , nil ) )
93+ count += len (encoding . Encode ( tool .Function .Arguments , nil , nil ) )
7294 }
73- count += len (msg .ToolCallID )
74- return count / 3
95+ count += len (encoding .Encode (msg .ToolCallID , nil , nil ))
96+
97+ return count , nil
98+ }
99+
100+ func countTools (tools []types.ChatCompletionTool ) (int , error ) {
101+ encoding , err := tiktoken .GetEncoding ("o200k_base" )
102+ if err != nil {
103+ return 0 , err
104+ }
105+
106+ toolJSON , err := json .Marshal (tools )
107+ if err != nil {
108+ return 0 , err
109+ }
110+
111+ return len (encoding .Encode (string (toolJSON ), nil , nil )), nil
75112}
0 commit comments