wehub-resource-sync
34 lines
1.2 KiB
Go
34 lines
1.2 KiB
Go
package middleware
|
|
|
|
import "github.com/labstack/echo/v4"
|
|
|
|
// StampUsage records the canonical token counts on the echo context so
|
|
// UsageMiddleware can attribute the request without parsing the response
|
|
// body. Handlers must call this for every successful response — the
|
|
// body-parse fallback is reserved for foreign endpoints (e.g., the cloud
|
|
// passthrough proxy).
|
|
//
|
|
// model is the name written into the response payload; passing it here
|
|
// is what lets the middleware fill the UsageRecord even when the handler
|
|
// abbreviates or rewrites the user-supplied model. Empty values are
|
|
// ignored so partial information is still useful (e.g., embeddings calls
|
|
// where completion is always 0).
|
|
//
|
|
// prompt and completion accept int because that's the native width of
|
|
// LocalAI's TokenUsage / OpenAIUsage structs (token counts never come
|
|
// close to overflow). Conversion to int64 happens once, here, so call
|
|
// sites stay free of casts.
|
|
func StampUsage(c echo.Context, model string, prompt, completion int) {
|
|
if c == nil {
|
|
return
|
|
}
|
|
if model != "" {
|
|
c.Set(ContextKeyResponseModel, model)
|
|
}
|
|
p := int64(prompt)
|
|
cp := int64(completion)
|
|
c.Set(ContextKeyPromptTokens, p)
|
|
c.Set(ContextKeyCompletionTokens, cp)
|
|
c.Set(ContextKeyTotalTokens, p+cp)
|
|
}
|