Skip to content

fix(usage): record real backend token counts and mark estimated rows #201

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
apapi merged 1 commit into from
Sep 28, 2026
Merged

fix(usage): record real backend token counts and mark estimated rows #201

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
139 changes: 76 additions & 63 deletions internal/config/api_auth.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -208,55 +208,58 @@ type APIUsageEvent struct {
LimitedCount int64
InputTokens int64
OutputTokens int64
Estimated bool
CreatedAt time.Time
}

type APIUsageRecord struct {
APIKeyID string `json:"api_key_id"`
APIKeyName string `json:"api_key_name"`
Model string `json:"model"`
Source string `json:"source,omitempty"`
SourceType string `json:"source_type,omitempty"`
SourceName string `json:"source_name,omitempty"`
PoolID string `json:"pool_id,omitempty"`
PoolName string `json:"pool_name,omitempty"`
PoolModel string `json:"pool_model,omitempty"`
ActualMemberID string `json:"actual_member_id,omitempty"`
MemberModel string `json:"member_model,omitempty"`
EstimatedCost float64 `json:"estimated_cost,omitempty"`
CostCurrency string `json:"cost_currency,omitempty"`
CostKnown bool `json:"cost_known"`
FallbackCount int64 `json:"fallback_count,omitempty"`
LimitedCount int64 `json:"limited_count,omitempty"`
Requests int64 `json:"requests"`
InputTokens int64 `json:"input_tokens"`
OutputTokens int64 `json:"output_tokens"`
TotalTokens int64 `json:"total_tokens"`
LastUsedAt time.Time `json:"last_used_at"`
APIKeyID string `json:"api_key_id"`
APIKeyName string `json:"api_key_name"`
Model string `json:"model"`
Source string `json:"source,omitempty"`
SourceType string `json:"source_type,omitempty"`
SourceName string `json:"source_name,omitempty"`
PoolID string `json:"pool_id,omitempty"`
PoolName string `json:"pool_name,omitempty"`
PoolModel string `json:"pool_model,omitempty"`
ActualMemberID string `json:"actual_member_id,omitempty"`
MemberModel string `json:"member_model,omitempty"`
EstimatedCost float64 `json:"estimated_cost,omitempty"`
CostCurrency string `json:"cost_currency,omitempty"`
CostKnown bool `json:"cost_known"`
FallbackCount int64 `json:"fallback_count,omitempty"`
LimitedCount int64 `json:"limited_count,omitempty"`
Requests int64 `json:"requests"`
EstimatedRequests int64 `json:"estimated_requests,omitempty"`
InputTokens int64 `json:"input_tokens"`
OutputTokens int64 `json:"output_tokens"`
TotalTokens int64 `json:"total_tokens"`
LastUsedAt time.Time `json:"last_used_at"`
}

type APIUsageEventRecord struct {
APIKeyID string `json:"api_key_id"`
APIKeyName string `json:"api_key_name"`
Model string `json:"model"`
Source string `json:"source,omitempty"`
SourceType string `json:"source_type,omitempty"`
SourceName string `json:"source_name,omitempty"`
PoolID string `json:"pool_id,omitempty"`
PoolName string `json:"pool_name,omitempty"`
PoolModel string `json:"pool_model,omitempty"`
ActualMemberID string `json:"actual_member_id,omitempty"`
MemberModel string `json:"member_model,omitempty"`
EstimatedCost float64 `json:"estimated_cost,omitempty"`
CostCurrency string `json:"cost_currency,omitempty"`
CostKnown bool `json:"cost_known"`
FallbackCount int64 `json:"fallback_count,omitempty"`
LimitedCount int64 `json:"limited_count,omitempty"`
Requests int64 `json:"requests,omitempty"`
InputTokens int64 `json:"input_tokens"`
OutputTokens int64 `json:"output_tokens"`
TotalTokens int64 `json:"total_tokens"`
CreatedAt time.Time `json:"created_at"`
APIKeyID string `json:"api_key_id"`
APIKeyName string `json:"api_key_name"`
Model string `json:"model"`
Source string `json:"source,omitempty"`
SourceType string `json:"source_type,omitempty"`
SourceName string `json:"source_name,omitempty"`
PoolID string `json:"pool_id,omitempty"`
PoolName string `json:"pool_name,omitempty"`
PoolModel string `json:"pool_model,omitempty"`
ActualMemberID string `json:"actual_member_id,omitempty"`
MemberModel string `json:"member_model,omitempty"`
EstimatedCost float64 `json:"estimated_cost,omitempty"`
CostCurrency string `json:"cost_currency,omitempty"`
CostKnown bool `json:"cost_known"`
FallbackCount int64 `json:"fallback_count,omitempty"`
LimitedCount int64 `json:"limited_count,omitempty"`
Requests int64 `json:"requests,omitempty"`
EstimatedRequests int64 `json:"estimated_requests,omitempty"`
InputTokens int64 `json:"input_tokens"`
OutputTokens int64 `json:"output_tokens"`
TotalTokens int64 `json:"total_tokens"`
CreatedAt time.Time `json:"created_at"`
}

type APIUsageState struct {
Expand Down Expand Up @@ -349,6 +352,7 @@ func compactAPIUsageEvents(events []APIUsageEventRecord) []APIUsageEventRecord {
out[i].SourceName = latestNonEmpty(out[i].SourceName, event.SourceName)
out[i].PoolName = latestNonEmpty(out[i].PoolName, event.PoolName)
out[i].Requests += event.Requests
out[i].EstimatedRequests += event.EstimatedRequests
out[i].FallbackCount += event.FallbackCount
out[i].LimitedCount += event.LimitedCount
out[i].EstimatedCost += event.EstimatedCost
Expand Down Expand Up @@ -384,6 +388,7 @@ func upsertAPIUsageRecord(state *APIUsageState, event APIUsageEventRecord) {
state.Records[i].SourceName = event.SourceName
state.Records[i].PoolName = latestNonEmpty(state.Records[i].PoolName, event.PoolName)
state.Records[i].Requests += requests
state.Records[i].EstimatedRequests += event.EstimatedRequests
state.Records[i].FallbackCount += event.FallbackCount
state.Records[i].LimitedCount += event.LimitedCount
state.Records[i].EstimatedCost += event.EstimatedCost
Expand All @@ -395,27 +400,28 @@ func upsertAPIUsageRecord(state *APIUsageState, event APIUsageEventRecord) {
}
}
state.Records = append(state.Records, APIUsageRecord{
APIKeyID: event.APIKeyID,
APIKeyName: event.APIKeyName,
Model: event.Model,
Source: event.Source,
SourceType: event.SourceType,
SourceName: event.SourceName,
PoolID: event.PoolID,
PoolName: event.PoolName,
PoolModel: event.PoolModel,
ActualMemberID: event.ActualMemberID,
MemberModel: event.MemberModel,
EstimatedCost: event.EstimatedCost,
CostCurrency: event.CostCurrency,
CostKnown: event.CostKnown,
FallbackCount: event.FallbackCount,
LimitedCount: event.LimitedCount,
Requests: requests,
InputTokens: event.InputTokens,
OutputTokens: event.OutputTokens,
TotalTokens: apiUsageEventTotalTokens(event),
LastUsedAt: event.CreatedAt,
APIKeyID: event.APIKeyID,
APIKeyName: event.APIKeyName,
Model: event.Model,
Source: event.Source,
SourceType: event.SourceType,
SourceName: event.SourceName,
PoolID: event.PoolID,
PoolName: event.PoolName,
PoolModel: event.PoolModel,
ActualMemberID: event.ActualMemberID,
MemberModel: event.MemberModel,
EstimatedCost: event.EstimatedCost,
CostCurrency: event.CostCurrency,
CostKnown: event.CostKnown,
FallbackCount: event.FallbackCount,
LimitedCount: event.LimitedCount,
Requests: requests,
EstimatedRequests: event.EstimatedRequests,
InputTokens: event.InputTokens,
OutputTokens: event.OutputTokens,
TotalTokens: apiUsageEventTotalTokens(event),
LastUsedAt: event.CreatedAt,
})
}

Expand Down Expand Up @@ -463,6 +469,13 @@ func apiUsageEventRequests(event APIUsageEventRecord) int64 {
return 1
}

func apiUsageEventEstimatedRequests(event APIUsageEventRecord) int64 {
if event.EstimatedRequests > 0 {
return event.EstimatedRequests
}
return 0
}

func apiUsageEventTotalTokens(event APIUsageEventRecord) int64 {
if event.TotalTokens != 0 {
return event.TotalTokens
Expand Down
101 changes: 76 additions & 25 deletions internal/config/api_usage_store.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@ CREATE TABLE IF NOT EXISTS api_usage_events (
fallback_count INTEGER NOT NULL DEFAULT 0,
limited_count INTEGER NOT NULL DEFAULT 0,
requests INTEGER NOT NULL DEFAULT 0,
estimated_requests INTEGER NOT NULL DEFAULT 0,
input_tokens INTEGER NOT NULL DEFAULT 0,
output_tokens INTEGER NOT NULL DEFAULT 0,
total_tokens INTEGER NOT NULL DEFAULT 0,
Expand All @@ -65,15 +66,15 @@ CREATE TABLE IF NOT EXISTS api_usage_meta (

const apiUsageColumns = `day, api_key_id, model, source, source_type, pool_id, pool_model,
actual_member_id, member_model, cost_currency, cost_known, api_key_name, source_name,
pool_name, estimated_cost, fallback_count, limited_count, requests, input_tokens,
output_tokens, total_tokens, created_at`
pool_name, estimated_cost, fallback_count, limited_count, requests, estimated_requests,
input_tokens, output_tokens, total_tokens, created_at`

// apiUsageUpsertStatement folds a new event into its day bucket in a single
// write. Names fall back to the stored value when the new event omits them,
// matching latestNonEmpty.
const apiUsageUpsertStatement = `
INSERT INTO api_usage_events (` + apiUsageColumns + `)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT (
day, api_key_id, model, source, source_type,
pool_id, pool_model, actual_member_id, member_model,
Expand All @@ -86,6 +87,7 @@ ON CONFLICT (
fallback_count = api_usage_events.fallback_count + excluded.fallback_count,
limited_count = api_usage_events.limited_count + excluded.limited_count,
requests = api_usage_events.requests + excluded.requests,
estimated_requests = api_usage_events.estimated_requests + excluded.estimated_requests,
input_tokens = api_usage_events.input_tokens + excluded.input_tokens,
output_tokens = api_usage_events.output_tokens + excluded.output_tokens,
total_tokens = api_usage_events.total_tokens + excluded.total_tokens,
Expand Down Expand Up @@ -186,6 +188,21 @@ func (s *APIUsageStore) openLocked() (*sql.DB, error) {
_ = db.Close()
return nil, fmt.Errorf("initializing API usage database: %w", err)
}
// Older databases created before the estimated_requests column existed
// need an ALTER TABLE; new databases already have it from the schema.
// Check PRAGMA table_info so we only ALTER when the column is missing,
// and any real failure surfaces instead of being silently swallowed.
needsMigration, err := needsEstimatedRequestsColumn(db)
if err != nil {
_ = db.Close()
return nil, fmt.Errorf("checking API usage database schema: %w", err)
}
if needsMigration {
if _, err := db.Exec("ALTER TABLE api_usage_events ADD COLUMN estimated_requests INTEGER NOT NULL DEFAULT 0"); err != nil {
_ = db.Close()
return nil, fmt.Errorf("migrating API usage database: %w", err)
}
}
s.db = db
// A failed import must not take usage recording down: the JSON file stays
// in place so the next start can retry it.
Expand Down Expand Up @@ -299,7 +316,7 @@ func apiUsageSelectEvents(db *sql.DB, options APIUsageListOptions) ([]APIUsageEv
&event.PoolID, &event.PoolModel, &event.ActualMemberID, &event.MemberModel,
&event.CostCurrency, &costKnown, &event.APIKeyName, &event.SourceName,
&event.PoolName, &event.EstimatedCost, &event.FallbackCount, &event.LimitedCount,
&event.Requests, &event.InputTokens, &event.OutputTokens, &event.TotalTokens,
&event.Requests, &event.EstimatedRequests, &event.InputTokens, &event.OutputTokens, &event.TotalTokens,
&createdAt,
); err != nil {
return nil, fmt.Errorf("reading API usage: %w", err)
Expand Down Expand Up @@ -347,27 +364,28 @@ func apiUsageEventRecord(event APIUsageEvent) (APIUsageEventRecord, bool) {
event.CostKnown, event.CostCurrency, event.EstimatedCost,
)
compacted := compactAPIUsageEvents([]APIUsageEventRecord{{
APIKeyID: event.APIKeyID,
APIKeyName: event.APIKeyName,
Model: event.Model,
Source: event.Source,
SourceType: event.SourceType,
SourceName: event.SourceName,
PoolID: event.PoolID,
PoolName: event.PoolName,
PoolModel: event.PoolModel,
ActualMemberID: event.ActualMemberID,
MemberModel: event.MemberModel,
EstimatedCost: estimatedCost,
CostCurrency: costCurrency,
CostKnown: costKnown,
FallbackCount: event.FallbackCount,
LimitedCount: event.LimitedCount,
Requests: 1,
InputTokens: event.InputTokens,
OutputTokens: event.OutputTokens,
TotalTokens: event.InputTokens + event.OutputTokens,
CreatedAt: createdAt,
APIKeyID: event.APIKeyID,
APIKeyName: event.APIKeyName,
Model: event.Model,
Source: event.Source,
SourceType: event.SourceType,
SourceName: event.SourceName,
PoolID: event.PoolID,
PoolName: event.PoolName,
PoolModel: event.PoolModel,
ActualMemberID: event.ActualMemberID,
MemberModel: event.MemberModel,
EstimatedCost: estimatedCost,
CostCurrency: costCurrency,
CostKnown: costKnown,
FallbackCount: event.FallbackCount,
LimitedCount: event.LimitedCount,
Requests: 1,
EstimatedRequests: boolToInt64(event.Estimated),
InputTokens: event.InputTokens,
OutputTokens: event.OutputTokens,
TotalTokens: event.InputTokens + event.OutputTokens,
CreatedAt: createdAt,
}})
if len(compacted) == 0 {
return APIUsageEventRecord{}, false
Expand Down Expand Up @@ -399,6 +417,7 @@ func apiUsageInsertArgs(event APIUsageEventRecord) []any {
event.FallbackCount,
event.LimitedCount,
apiUsageEventRequests(event),
apiUsageEventEstimatedRequests(event),
event.InputTokens,
event.OutputTokens,
apiUsageEventTotalTokens(event),
Expand All @@ -413,9 +432,41 @@ func apiUsageTimeToStorage(value time.Time) int64 {
return value.UTC().UnixNano()
}

func boolToInt64(v bool) int64 {
if v {
return 1
}
return 0
}

func apiUsageTimeFromStorage(value int64) time.Time {
if value == 0 {
return time.Time{}
}
return time.Unix(0, value).UTC()
}

func needsEstimatedRequestsColumn(db *sql.DB) (bool, error) {
rows, err := db.Query("PRAGMA table_info(api_usage_events)")
if err != nil {
return false, fmt.Errorf("querying table_info: %w", err)
}
defer rows.Close()
for rows.Next() {
var cid int
var name, ctype string
var notnull int
var dfltValue any
var pk int
if err := rows.Scan(&cid, &name, &ctype, &notnull, &dfltValue, &pk); err != nil {
return false, fmt.Errorf("scanning table_info: %w", err)
}
if name == "estimated_requests" {
return false, nil
}
}
if err := rows.Err(); err != nil {
return false, fmt.Errorf("reading table_info: %w", err)
}
return true, nil
}
16 changes: 16 additions & 0 deletions internal/inference/llama.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1026,6 +1026,7 @@ func trimOldestNonSystemMessage(messages []Message) ([]Message, bool) {

type llamaChatUsage struct {
PromptTokens int64 `json:"prompt_tokens"`
CompletionTokens int64 `json:"completion_tokens"`
PromptTokensDetails struct {
CachedTokens int64 `json:"cached_tokens"`
} `json:"prompt_tokens_details"`
Expand Down Expand Up @@ -1057,6 +1058,19 @@ func recordLlamaCacheUsage(collector *CacheUsageCollector, usage llamaChatUsage,
}
}

// reportLlamaGenerationUsage forwards the prompt/completion token counts a
// backend advertised to the caller's OnUsage callback. It is a no-op when the
// caller did not request usage reporting or the backend reported nothing.
func reportLlamaGenerationUsage(onUsage func(int64, int64), usage llamaChatUsage) {
if onUsage == nil {
return
}
if usage.PromptTokens <= 0 && usage.CompletionTokens <= 0 {
return
}
onUsage(usage.PromptTokens, usage.CompletionTokens)
}

func (e *llamaEngine) handleStream(body io.Reader, onToken TokenCallback, opts Options) (string, error) {
return e.handleStreamWithCacheUsage(body, onToken, opts, nil)
}
Expand Down Expand Up @@ -1092,6 +1106,7 @@ func (e *llamaEngine) handleStreamWithCacheUsage(body io.Reader, onToken TokenCa
continue
}
recordLlamaCacheUsage(collector, chunk.Usage, chunk.Timings)
reportLlamaGenerationUsage(opts.OnUsage, chunk.Usage)
if len(chunk.Choices) > 0 {
d := chunk.Choices[0].Delta
// Use at most one delta text per chunk. Some llama-server builds populate both
Expand Down Expand Up @@ -1146,6 +1161,7 @@ func (e *llamaEngine) handleNonStreamWithCacheUsage(body io.Reader, opts Options
return "", fmt.Errorf("decoding response: %w", err)
}
recordLlamaCacheUsage(collector, resp.Usage, resp.Timings)
reportLlamaGenerationUsage(opts.OnUsage, resp.Usage)
if len(resp.Choices) == 0 {
return "", fmt.Errorf("no choices in response")
}
Expand Down
Loading
Loading