| 18 | } |
| 19 | |
| 20 | func AddUsage(state *AgentState, usageOrInput any, outputTokens ...int) { |
| 21 | switch usage := usageOrInput.(type) { |
| 22 | case map[string]int: |
| 23 | state.TotalInputTokens += usage["input_tokens"] |
| 24 | state.TotalOutputTokens += usage["output_tokens"] |
| 25 | state.CacheReadInputTokens += usage["cache_read_input_tokens"] |
| 26 | state.CacheCreateInputTokens += usage["cache_creation_input_tokens"] |
| 27 | state.ServerToolUseInputTokens += usage["server_tool_use_input_tokens"] |
| 28 | case map[string]any: |
| 29 | state.TotalInputTokens += intFromAny(usage["input_tokens"]) |
| 30 | state.TotalOutputTokens += intFromAny(usage["output_tokens"]) |
| 31 | state.CacheReadInputTokens += intFromAny(usage["cache_read_input_tokens"]) |
| 32 | state.CacheCreateInputTokens += intFromAny(usage["cache_creation_input_tokens"]) |
| 33 | state.ServerToolUseInputTokens += intFromAny(usage["server_tool_use_input_tokens"]) |
| 34 | case int: |
| 35 | state.TotalInputTokens += usage |
| 36 | if len(outputTokens) > 0 { |
| 37 | state.TotalOutputTokens += outputTokens[0] |
| 38 | } |
| 39 | } |
| 40 | } |
| 41 | |
| 42 | func CommitAssistantTurn(state *AgentState, thinkingContent []map[string]any, fullText string, toolCalls []coretools.ToolCall, messageID string, inputTokens, outputTokens int) { |
| 43 | AddUsage(state, inputTokens, outputTokens) |