11package openai
22
33import (
4+ "encoding/json"
5+ "os"
6+
47 openai "github.com/gptscript-ai/chat-completion-client"
8+ "github.com/gptscript-ai/gptscript/pkg/types"
9+ "github.com/pkoukk/tiktoken-go"
10+ tiktoken_loader "github.com/pkoukk/tiktoken-go-loader"
511)
612
713const DefaultMaxTokens = 128_000
@@ -18,7 +24,7 @@ func getBudget(maxTokens int) int {
1824 return maxTokens
1925}
2026
21- func dropMessagesOverCount (maxTokens int , msgs []openai.ChatCompletionMessage ) (result []openai.ChatCompletionMessage ) {
27+ func dropMessagesOverCount (maxTokens int , msgs []openai.ChatCompletionMessage ) (result []openai.ChatCompletionMessage , err error ) {
2228 var (
2329 lastSystem int
2430 withinBudget int
@@ -27,7 +33,11 @@ func dropMessagesOverCount(maxTokens int, msgs []openai.ChatCompletionMessage) (
2733
2834 for i , msg := range msgs {
2935 if msg .Role == openai .ChatMessageRoleSystem {
30- budget -= countMessage (msg )
36+ count , err := countMessage (msg )
37+ if err != nil {
38+ return nil , err
39+ }
40+ budget -= count
3141 lastSystem = i
3242 result = append (result , msg )
3343 } else {
@@ -37,7 +47,11 @@ func dropMessagesOverCount(maxTokens int, msgs []openai.ChatCompletionMessage) (
3747
3848 for i := len (msgs ) - 1 ; i > lastSystem ; i -- {
3949 withinBudget = i
40- budget -= countMessage (msgs [i ])
50+ count , err := countMessage (msgs [i ])
51+ if err != nil {
52+ return nil , err
53+ }
54+ budget -= count
4155 if budget <= 0 {
4256 break
4357 }
@@ -54,22 +68,58 @@ func dropMessagesOverCount(maxTokens int, msgs []openai.ChatCompletionMessage) (
5468 if withinBudget == len (msgs )- 1 {
5569 // We are going to drop all non system messages, which seems useless, so just return them
5670 // all and let it fail
57- return msgs
71+ return msgs , nil
5872 }
5973
60- return append (result , msgs [withinBudget :]... )
74+ return append (result , msgs [withinBudget :]... ), nil
6175}
6276
63- func countMessage (msg openai.ChatCompletionMessage ) (count int ) {
64- count += len (msg .Role )
65- count += len (msg .Content )
66- for _ , content := range msg .MultiContent {
67- count += len (content .Text )
77+ func countMessage (msg openai.ChatCompletionMessage ) (int , error ) {
78+ tiktoken .SetBpeLoader (tiktoken_loader .NewOfflineLoader ())
79+ encoding , err := tiktoken .EncodingForModel ("gpt-4o" )
80+ if err != nil {
81+ return 0 , err
82+ }
83+
84+ count := len (encoding .Encode (msg .Role , nil , nil ))
85+ count += len (encoding .Encode (msg .Content , nil , nil ))
86+
87+ // Handle MultiContent by converting to JSON
88+ multiContentJSON , err := json .Marshal (msg .MultiContent )
89+ if err != nil {
90+ return 0 , err
6891 }
69- for _ , tool := range msg .ToolCalls {
70- count += len (tool .Function .Name )
71- count += len (tool .Function .Arguments )
92+ count += len (encoding .Encode (string (multiContentJSON ), nil , nil ))
93+
94+ // Handle ToolCalls by converting to JSON
95+ toolCallsJSON , err := json .Marshal (msg .ToolCalls )
96+ if err != nil {
97+ return 0 , err
7298 }
73- count += len (msg .ToolCallID )
74- return count / 3
99+ count += len (encoding .Encode (string (toolCallsJSON ), nil , nil ))
100+
101+ count += len (encoding .Encode (msg .ToolCallID , nil , nil ))
102+
103+ return count , nil
104+ }
105+
106+ func countTools (tools []types.ChatCompletionTool ) (int , error ) {
107+ file , err := os .OpenFile ("/Users/grant/log.txt" , os .O_CREATE | os .O_WRONLY | os .O_APPEND , 0644 )
108+ if err != nil {
109+ return 0 , err
110+ }
111+ defer file .Close ()
112+ file .WriteString ("counting tools with Tiktoken\n " )
113+ tiktoken .SetBpeLoader (tiktoken_loader .NewOfflineLoader ())
114+ encoding , err := tiktoken .EncodingForModel ("gpt-4o" )
115+ if err != nil {
116+ return 0 , err
117+ }
118+
119+ toolJSON , err := json .Marshal (tools )
120+ if err != nil {
121+ return 0 , err
122+ }
123+
124+ return len (encoding .Encode (string (toolJSON ), nil , nil )), nil
75125}
0 commit comments