1package agent
2
3import (
4 "fmt"
5
6 "charm.land/fantasy"
7)
8
9func usageIsZero(usage fantasy.Usage) bool {
10 return usage.InputTokens == 0 &&
11 usage.OutputTokens == 0 &&
12 usage.TotalTokens == 0 &&
13 usage.ReasoningTokens == 0 &&
14 usage.CacheCreationTokens == 0 &&
15 usage.CacheReadTokens == 0
16}
17
18func fallbackStepUsage(messages []fantasy.Message, step fantasy.StepResult) (fantasy.Usage, bool) {
19 if !usageIsZero(step.Usage) {
20 return step.Usage, false
21 }
22
23 inputTokens := estimateMessageTokens(messages)
24 outputTokens := estimateStepCompletionTokens(step)
25 if inputTokens == 0 && outputTokens == 0 {
26 return fantasy.Usage{}, false
27 }
28
29 return fantasy.Usage{
30 InputTokens: inputTokens,
31 OutputTokens: outputTokens,
32 TotalTokens: inputTokens + outputTokens,
33 }, true
34}
35
36func cloneFantasyMessages(messages []fantasy.Message) []fantasy.Message {
37 cloned := make([]fantasy.Message, len(messages))
38 for i, msg := range messages {
39 cloned[i] = msg
40 cloned[i].Content = append([]fantasy.MessagePart(nil), msg.Content...)
41 }
42 return cloned
43}
44
45func estimateMessageTokens(messages []fantasy.Message) int64 {
46 var tokens int64
47 for _, msg := range messages {
48 tokens += approxTokenCount(string(msg.Role))
49 for _, part := range msg.Content {
50 tokens += estimateMessagePartTokens(part)
51 }
52 }
53 return tokens
54}
55
56func estimateStepCompletionTokens(step fantasy.StepResult) int64 {
57 var tokens int64
58 for _, content := range step.Content {
59 switch c := content.(type) {
60 case fantasy.TextContent:
61 tokens += approxTokenCount(c.Text)
62 case *fantasy.TextContent:
63 tokens += approxTokenCount(c.Text)
64 case fantasy.ReasoningContent:
65 tokens += approxTokenCount(c.Text)
66 case *fantasy.ReasoningContent:
67 tokens += approxTokenCount(c.Text)
68 case fantasy.FileContent:
69 tokens += estimateGeneratedFileTokens(c)
70 case *fantasy.FileContent:
71 tokens += estimateGeneratedFileTokens(*c)
72 case fantasy.SourceContent:
73 tokens += estimateSourceTokens(c)
74 case *fantasy.SourceContent:
75 tokens += estimateSourceTokens(*c)
76 case fantasy.ToolCallContent:
77 tokens += estimateToolCallTokens(c.ToolName, c.Input)
78 case *fantasy.ToolCallContent:
79 tokens += estimateToolCallTokens(c.ToolName, c.Input)
80 case fantasy.ToolResultContent:
81 if c.ProviderExecuted {
82 tokens += estimateToolResultContentTokens(c.ToolCallID, c.ToolName, c.ClientMetadata, c.Result)
83 }
84 case *fantasy.ToolResultContent:
85 if c.ProviderExecuted {
86 tokens += estimateToolResultContentTokens(c.ToolCallID, c.ToolName, c.ClientMetadata, c.Result)
87 }
88 }
89 }
90 return tokens
91}
92
93func estimateMessagePartTokens(part fantasy.MessagePart) int64 {
94 switch p := part.(type) {
95 case fantasy.TextPart:
96 return approxTokenCount(p.Text)
97 case *fantasy.TextPart:
98 return approxTokenCount(p.Text)
99 case fantasy.ReasoningPart:
100 return approxTokenCount(p.Text)
101 case *fantasy.ReasoningPart:
102 return approxTokenCount(p.Text)
103 case fantasy.FilePart:
104 return estimateFilePartTokens(p)
105 case *fantasy.FilePart:
106 return estimateFilePartTokens(*p)
107 case fantasy.ToolCallPart:
108 return estimateToolCallTokens(p.ToolName, p.Input)
109 case *fantasy.ToolCallPart:
110 return estimateToolCallTokens(p.ToolName, p.Input)
111 case fantasy.ToolResultPart:
112 return estimateToolResultContentTokens(p.ToolCallID, "", "", p.Output)
113 case *fantasy.ToolResultPart:
114 return estimateToolResultContentTokens(p.ToolCallID, "", "", p.Output)
115 default:
116 return 0
117 }
118}
119
120func estimateToolCallTokens(toolName, input string) int64 {
121 return approxTokenCount(toolName) + approxTokenCount(input)
122}
123
124func estimateToolResultContentTokens(toolCallID, toolName, metadata string, output fantasy.ToolResultOutputContent) int64 {
125 tokens := approxTokenCount(toolCallID) + approxTokenCount(toolName) + approxTokenCount(metadata)
126 switch result := output.(type) {
127 case fantasy.ToolResultOutputContentText:
128 tokens += approxTokenCount(result.Text)
129 case *fantasy.ToolResultOutputContentText:
130 tokens += approxTokenCount(result.Text)
131 case fantasy.ToolResultOutputContentError:
132 if result.Error != nil {
133 tokens += approxTokenCount(result.Error.Error())
134 }
135 case *fantasy.ToolResultOutputContentError:
136 if result.Error != nil {
137 tokens += approxTokenCount(result.Error.Error())
138 }
139 case fantasy.ToolResultOutputContentMedia:
140 tokens += estimateMediaTokens(result.MediaType, result.Text, len(result.Data))
141 case *fantasy.ToolResultOutputContentMedia:
142 tokens += estimateMediaTokens(result.MediaType, result.Text, len(result.Data))
143 }
144 return tokens
145}
146
147func estimateFilePartTokens(file fantasy.FilePart) int64 {
148 return estimateMediaTokens(file.MediaType, file.Filename, len(file.Data))
149}
150
151func estimateGeneratedFileTokens(file fantasy.FileContent) int64 {
152 return estimateMediaTokens(file.MediaType, "", len(file.Data))
153}
154
155func estimateMediaTokens(mediaType, text string, dataBytes int) int64 {
156 if dataBytes == 0 {
157 return approxTokenCount(mediaType) + approxTokenCount(text)
158 }
159 return approxTokenCount(fmt.Sprintf("%s %s %d bytes", mediaType, text, dataBytes))
160}
161
162func estimateSourceTokens(source fantasy.SourceContent) int64 {
163 return approxTokenCount(string(source.SourceType)) +
164 approxTokenCount(source.ID) +
165 approxTokenCount(source.URL) +
166 approxTokenCount(source.Title) +
167 approxTokenCount(source.MediaType) +
168 approxTokenCount(source.Filename)
169}
170
171func approxTokenCount(s string) int64 {
172 if s == "" {
173 return 0
174 }
175 return int64((len(s) + 3) / 4)
176}