MCPcopy
hub / github.com/mudler/LocalAI / tokensFromBody

Function tokensFromBody

core/http/middleware/usage.go:181–238  ·  view source on GitHub ↗

tokensFromBody covers the passthrough-proxy / foreign-endpoint case where no handler stamps the context. Returns ok=false on any parse failure or missing-usage; the caller increments the unrecorded counter.

(responseBytes []byte, contentType string)

Source from the content-addressed store, hash-verified

179// where no handler stamps the context. Returns ok=false on any parse
180// failure or missing-usage; the caller increments the unrecorded counter.
181func tokensFromBody(responseBytes []byte, contentType string) (model string, prompt, completion, total int64, ok bool) {
182 if len(responseBytes) == 0 {
183 return
184 }
185 isJSON := contentType == "" || contentType == "application/json" || bytes.HasPrefix([]byte(contentType), []byte("application/json"))
186 isSSE := bytes.HasPrefix([]byte(contentType), []byte("text/event-stream"))
187 if !isJSON && !isSSE {
188 return
189 }
190
191 payload := responseBytes
192 if isSSE {
193 // For SSE, the canonical usage chunk is the *last* non-[DONE] data
194 // line. OpenAI clients only emit one if stream_options.include_usage
195 // is set; Anthropic emits a final message_delta with usage. Both
196 // fit the "last data: line" rule.
197 last, lastOk := lastSSEData(responseBytes)
198 if !lastOk {
199 return
200 }
201 payload = last
202 }
203
204 // Try OpenAI shape first (handles /v1/chat/completions, /v1/completions,
205 // /v1/embeddings, /v1/edits, and any proxy that translates to OpenAI).
206 // A usage block whose token fields all decoded to zero is ambiguous —
207 // it could be an Anthropic body that happens to have a `usage` key —
208 // so fall through to the Anthropic parser instead of recording zeros.
209 var openAI usageResponseBody
210 if err := json.Unmarshal(payload, &openAI); err == nil && openAI.Usage != nil {
211 if openAI.Usage.PromptTokens != 0 || openAI.Usage.CompletionTokens != 0 || openAI.Usage.TotalTokens != 0 {
212 model = openAI.Model
213 prompt = openAI.Usage.PromptTokens
214 completion = openAI.Usage.CompletionTokens
215 total = openAI.Usage.TotalTokens
216 if total == 0 {
217 total = prompt + completion
218 }
219 ok = true
220 return
221 }
222 }
223
224 // Fall through to Anthropic shape (proxy passthrough territory).
225 var ant anthropicResponseBody
226 if err := json.Unmarshal(payload, &ant); err == nil && ant.Usage != nil {
227 if ant.Usage.InputTokens != 0 || ant.Usage.OutputTokens != 0 {
228 model = ant.Model
229 prompt = ant.Usage.InputTokens
230 completion = ant.Usage.OutputTokens
231 total = prompt + completion
232 ok = true
233 return
234 }
235 }
236
237 return
238}

Callers 1

UsageMiddlewareFunction · 0.85

Calls 1

lastSSEDataFunction · 0.85

Tested by

no test coverage detected