mirror of
https://github.com/mudler/LocalAI.git
synced 2026-09-22 14:14:54 -04:00
Richard Palethorpe
495165266f
Report input tokens and frame-step output units in response metadata and record them through the existing usage accounting pipeline. Preserve the accounting rule and model-specific dimensions as JSON without extending the gRPC schema for each modality. Expose animation usage only under metadata.usage, validate counts before recording, and document the response contract and loaded-model location. Add coverage for transport, defaults, failures, persistence, and recording requests once with statistics enabled or disabled. Assisted-by: Codex:GPT-6 Signed-off-by: Richard Palethorpe <io@richiejp.com>
55 lines
1.8 KiB
Go
55 lines
1.8 KiB
Go
package middleware
|
|
|
|
import (
|
|
"github.com/labstack/echo/v4"
|
|
"github.com/mudler/LocalAI/pkg/grpc/metadata"
|
|
)
|
|
|
|
const responseMetadataKey = "localai_response_metadata"
|
|
|
|
// StampResponseMetadata extracts the shared usage convention while preserving
|
|
// arbitrary metadata. Invalid counts never enter the accounting pipeline.
|
|
func StampResponseMetadata(c echo.Context, model string, data []byte) error {
|
|
usage, err := metadata.ParseUsage(data)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if c != nil {
|
|
c.Set(responseMetadataKey, string(data))
|
|
if usage != nil {
|
|
StampUsage(c, model, usage.InputUnits, usage.OutputUnits)
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// 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)
|
|
}
|