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 tokens += estimateToolResultContentTokens(c.ToolCallID, c.ToolName, c.ClientMetadata, c.Result)
82 case *fantasy.ToolResultContent:
83 tokens += estimateToolResultContentTokens(c.ToolCallID, c.ToolName, c.ClientMetadata, c.Result)
84 }
85 }
86 return tokens
87}
88
89func estimateMessagePartTokens(part fantasy.MessagePart) int64 {
90 switch p := part.(type) {
91 case fantasy.TextPart:
92 return approxTokenCount(p.Text)
93 case *fantasy.TextPart:
94 return approxTokenCount(p.Text)
95 case fantasy.ReasoningPart:
96 return approxTokenCount(p.Text)
97 case *fantasy.ReasoningPart:
98 return approxTokenCount(p.Text)
99 case fantasy.FilePart:
100 return estimateFilePartTokens(p)
101 case *fantasy.FilePart:
102 return estimateFilePartTokens(*p)
103 case fantasy.ToolCallPart:
104 return estimateToolCallTokens(p.ToolName, p.Input)
105 case *fantasy.ToolCallPart:
106 return estimateToolCallTokens(p.ToolName, p.Input)
107 case fantasy.ToolResultPart:
108 return estimateToolResultContentTokens(p.ToolCallID, "", "", p.Output)
109 case *fantasy.ToolResultPart:
110 return estimateToolResultContentTokens(p.ToolCallID, "", "", p.Output)
111 default:
112 return 0
113 }
114}
115
116func estimateToolCallTokens(toolName, input string) int64 {
117 return approxTokenCount(toolName) + approxTokenCount(input)
118}
119
120func estimateToolResultContentTokens(toolCallID, toolName, metadata string, output fantasy.ToolResultOutputContent) int64 {
121 tokens := approxTokenCount(toolCallID) + approxTokenCount(toolName) + approxTokenCount(metadata)
122 switch result := output.(type) {
123 case fantasy.ToolResultOutputContentText:
124 tokens += approxTokenCount(result.Text)
125 case *fantasy.ToolResultOutputContentText:
126 tokens += approxTokenCount(result.Text)
127 case fantasy.ToolResultOutputContentError:
128 if result.Error != nil {
129 tokens += approxTokenCount(result.Error.Error())
130 }
131 case *fantasy.ToolResultOutputContentError:
132 if result.Error != nil {
133 tokens += approxTokenCount(result.Error.Error())
134 }
135 case fantasy.ToolResultOutputContentMedia:
136 tokens += estimateMediaTokens(result.MediaType, result.Text, len(result.Data))
137 case *fantasy.ToolResultOutputContentMedia:
138 tokens += estimateMediaTokens(result.MediaType, result.Text, len(result.Data))
139 }
140 return tokens
141}
142
143func estimateFilePartTokens(file fantasy.FilePart) int64 {
144 return estimateMediaTokens(file.MediaType, file.Filename, len(file.Data))
145}
146
147func estimateGeneratedFileTokens(file fantasy.FileContent) int64 {
148 return estimateMediaTokens(file.MediaType, "", len(file.Data))
149}
150
151func estimateMediaTokens(mediaType, text string, dataBytes int) int64 {
152 if dataBytes == 0 {
153 return approxTokenCount(mediaType) + approxTokenCount(text)
154 }
155 return approxTokenCount(fmt.Sprintf("%s %s %d bytes", mediaType, text, dataBytes))
156}
157
158func estimateSourceTokens(source fantasy.SourceContent) int64 {
159 return approxTokenCount(string(source.SourceType)) +
160 approxTokenCount(source.ID) +
161 approxTokenCount(source.URL) +
162 approxTokenCount(source.Title) +
163 approxTokenCount(source.MediaType) +
164 approxTokenCount(source.Filename)
165}
166
167func approxTokenCount(s string) int64 {
168 if s == "" {
169 return 0
170 }
171 return int64((len(s) + 3) / 4)
172}