| 109 | } |
| 110 | |
| 111 | func TokenCountWithEstimation(messages []map[string]any) int { |
| 112 | if len(messages) == 0 { |
| 113 | return 0 |
| 114 | } |
| 115 | |
| 116 | n := len(messages) |
| 117 | |
| 118 | // Step 1: Find the last anchor: assistant message with usage.input_tokens. |
| 119 | anchorIdx := -1 |
| 120 | |
| 121 | for i := n - 1; i >= 0; i-- { |
| 122 | msg := messages[i] |
| 123 | |
| 124 | role, _ := msg["role"].(string) |
| 125 | if role != "assistant" { |
| 126 | continue |
| 127 | } |
| 128 | |
| 129 | usage, ok := msg["usage"].(map[string]any) |
| 130 | if !ok { |
| 131 | continue |
| 132 | } |
| 133 | |
| 134 | if _, exists := usage["input_tokens"]; exists { |
| 135 | anchorIdx = i |
| 136 | break |
| 137 | } |
| 138 | } |
| 139 | |
| 140 | if anchorIdx == -1 { |
| 141 | return RoughEstimate(messages) |
| 142 | } |
| 143 | |
| 144 | // Step 2: Extend anchor backwards for split responses with same id. |
| 145 | anchorMsg := messages[anchorIdx] |
| 146 | anchorID, hasAnchorID := anchorMsg["id"] |
| 147 | |
| 148 | if hasAnchorID && anchorID != nil { |
| 149 | for anchorIdx > 0 { |
| 150 | prev := messages[anchorIdx-1] |
| 151 | |
| 152 | prevID, ok := prev["id"] |
| 153 | if !ok || prevID != anchorID { |
| 154 | break |
| 155 | } |
| 156 | |
| 157 | anchorIdx-- |
| 158 | } |
| 159 | } |
| 160 | |
| 161 | // Step 3: Exact anchor count. |
| 162 | usage, ok := messages[anchorIdx]["usage"].(map[string]any) |
| 163 | if !ok { |
| 164 | usage = map[string]any{} |
| 165 | } |
| 166 | |
| 167 | // Walk forward within the split group to find the message that actually has usage. |
| 168 | for j := anchorIdx; j < n; j++ { |