usage

package
v0.1.56 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Jul 21, 2026 License: MIT Imports: 34 Imported by: 0

Documentation

Overview

Package usage provides token usage tracking for the AI gateway. It captures detailed token usage from API responses and stores them for analytics.

Index

Constants

View Source 
const (
	CacheTypeExact    = "exact"
	CacheTypeSemantic = "semantic"

	CacheModeUncached = "uncached"
	CacheModeCached   = "cached"
	CacheModeAll      = "all"
)
View Source 
const (
	// UsageEntryKey is the context key for storing the usage entry.
	UsageEntryKey contextKey = "usage_entry"

	// UsageEntryStreamingKey is the context key for marking a request as streaming.
	// When true, the middleware skips logging because streaming usage is handled
	// by the shared SSE observer path.
	UsageEntryStreamingKey contextKey = "usage_entry_streaming"
)
View Source 
const (
	CostSourceModelPricing      = "model_pricing"
	CostSourceOpenRouterCredits = "openrouter_credits"
	CostSourceXAITicks          = "xai_cost_in_usd_ticks"
)
View Source 
const (
	// DefaultDateRangeDays is the usage window applied when a query gives no
	// explicit range.
	DefaultDateRangeDays = 30
	// MaxDateRangeDays caps a requested usage window.
	MaxDateRangeDays = 365
)
View Source 
const (
	LiveEventUsageCompleted = "usage.completed"
	LiveEventUsageFailed    = "usage.failed"
	LiveEventUsageFlushed   = "usage.flushed"
)
View Source 
const (
	// BatchFlushThreshold is the number of entries that triggers an immediate flush.
	// When the batch reaches this size, it's written to storage without waiting for the timer.
	BatchFlushThreshold = 100
)

Buffer and batch limits for usage tracking.

View Source 
const CleanupInterval = 1 * time.Hour

CleanupInterval is how often the cleanup goroutine runs to delete old usage entries.

Variables

View Source 
var ErrPartialWrite = errors.New("partial write failure")

ErrPartialWrite indicates that a batch write only partially succeeded. Use errors.As to extract details about the failure.

Functions

func ApplyRewriteSavings 

func ApplyRewriteSavings(entry *UsageEntry, tokensSaved int, pricing *core.ModelPricing)

ApplyRewriteSavings folds a request-rewrite savings estimate into a usage entry: RewriteTokensSaved always, and RewriteCostSaved when pricing allows costing the removed tokens.

func BuildDateRange 

func BuildDateRange(startStr, endStr string, days int, location *time.Location, today time.Time) (time.Time, time.Time, error)

BuildDateRange resolves an inclusive [start, end] day range from optional YYYY-MM-DD strings. When only one bound is given the other defaults (start: 30 days before end, end: today); when neither is given the range covers the last days days ending today.

func EnrichUsageLogEntry 

func EnrichUsageLogEntry(entry *UsageLogEntry)

EnrichUsageLogEntry populates the derived provider-cache fields on entry (uncached/cached/write split and the cached ratio) from RawData. Safe to call on entries whose RawData does not contain provider cache fields.

func EntryInputSegments 

func EntryInputSegments(entry UsageLogEntry) (uncachedInput, cachedInput, cacheWriteInput int64)

EntryInputSegments splits one usage log entry's input tokens into the provider-uncached prompt, the provider-cached read, and the provider cache write portions. Provider-specific quirks are handled here so callers — request summaries, the admin usage log, and the live SSE preview — stay in sync. The various provider field names are coalesced via max:

  • cached reads: cache_read_input_tokens (Anthropic, Bedrock), prompt_cached_tokens (OpenAI, DeepSeek), cached_tokens (Gemini)
  • cache writes: cache_creation_input_tokens (Anthropic), cache_write_input_tokens (Bedrock Converse)

func NormalizeDateRangeDays 

func NormalizeDateRangeDays(days int) int

NormalizeDateRangeDays clamps days to [1, MaxDateRangeDays], defaulting to DefaultDateRangeDays when not positive.

func SummarizeUsageByRequestID 

func SummarizeUsageByRequestID(entriesByRequest map[string][]UsageLogEntry) map[string]*RequestUsageSummary

SummarizeUsageByRequestID aggregates usage log entries for each request ID.

func SummarizeUsageForRequestIDs 

func SummarizeUsageForRequestIDs(ctx context.Context, reader RequestUsageLoader, requestIDs []string) (map[string]*RequestUsageSummary, error)

SummarizeUsageForRequestIDs loads usage entries for requestIDs from reader and returns per-request summaries keyed by request ID. A nil reader or an empty ID list yields nil with no error.

Types

type CacheOverview 

type CacheOverview struct {
	Summary CacheOverviewSummary `json:"summary"`
	Daily   []CacheOverviewDaily `json:"daily"`
}

CacheOverview aggregates cached-only summary and daily series for the dashboard.

type CacheOverviewDaily 

type CacheOverviewDaily struct {
	Date         string   `json:"date"`
	Hits         int      `json:"hits"`
	ExactHits    int      `json:"exact_hits"`
	SemanticHits int      `json:"semantic_hits"`
	InputTokens  int64    `json:"input_tokens"`
	OutputTokens int64    `json:"output_tokens"`
	TotalTokens  int64    `json:"total_tokens"`
	SavedCost    *float64 `json:"saved_cost"`
}

CacheOverviewDaily holds cached-only statistics for a single period.

type CacheOverviewSummary 

type CacheOverviewSummary struct {
	TotalHits      int      `json:"total_hits"`
	ExactHits      int      `json:"exact_hits"`
	SemanticHits   int      `json:"semantic_hits"`
	TotalInput     int64    `json:"total_input_tokens"`
	TotalOutput    int64    `json:"total_output_tokens"`
	TotalTokens    int64    `json:"total_tokens"`
	TotalSavedCost *float64 `json:"total_saved_cost"`
}

CacheOverviewSummary holds cached-only aggregate statistics over a time period.

type Config 

type Config struct {
	// Enabled controls whether usage tracking is active
	Enabled bool

	// EnforceReturningUsageData controls whether to ask streaming providers to return usage data when possible.
	// When true, stream_options: {"include_usage": true} is added for provider paths that support it.
	// Default: true
	EnforceReturningUsageData bool

	// BufferSize is the number of usage entries to buffer before flushing
	BufferSize int

	// FlushInterval is how often to flush buffered entries
	FlushInterval time.Duration

	// RetentionDays is how long to keep usage data (0 = forever)
	RetentionDays int
}

Config holds usage tracking configuration

func DefaultConfig 

func DefaultConfig() Config

DefaultConfig returns a Config with sensible defaults

type CostResult 

type CostResult struct {
	InputCost  *float64
	OutputCost *float64
	TotalCost  *float64
	Caveat     string
	Source     string
}

CostResult holds the result of a granular cost calculation.

func CalculateGranularCost 

func CalculateGranularCost(inputTokens, outputTokens int, rawData map[string]any, providerType string, pricing *core.ModelPricing) CostResult

CalculateGranularCost computes input, output, and total costs from token counts, raw provider-specific data, and pricing information. It accounts for cached tokens, reasoning tokens, audio tokens, and other provider-specific token types.

The caveat field in the result describes any unmapped token fields or missing pricing data that prevented full cost calculation.

func CalculateUsageCost 

func CalculateUsageCost(inputTokens, outputTokens int, rawData map[string]any, providerType string, pricing *core.ModelPricing) CostResult

CalculateUsageCost prefers provider-supplied exact costs when available and falls back to static model pricing otherwise.

type DailyUsage 

type DailyUsage struct {
	Date         string `json:"date"`
	Requests     int    `json:"requests"`
	InputTokens  int64  `json:"input_tokens"`
	OutputTokens int64  `json:"output_tokens"`
	TotalTokens  int64  `json:"total_tokens"`
	// Provider prompt-cache split of the period's input, folded per row from
	// raw_data (same source as the summary). Zero when the storage layer does
	// not populate them.
	UncachedInputTokens   int64    `json:"uncached_input_tokens,omitempty"`
	CachedInputTokens     int64    `json:"cached_input_tokens,omitempty"`
	CacheWriteInputTokens int64    `json:"cache_write_input_tokens,omitempty"`
	InputCost             *float64 `json:"input_cost"`
	OutputCost            *float64 `json:"output_cost"`
	TotalCost             *float64 `json:"total_cost"`
}

DailyUsage holds usage statistics for a single period. Date holds the period label: YYYY-MM-DD for daily, YYYY-Www for weekly, YYYY-MM for monthly, or YYYY for yearly intervals.

type LabelUsage 

type LabelUsage struct {
	Label        string   `json:"label"`
	Requests     int      `json:"requests"`
	InputTokens  int64    `json:"input_tokens"`
	OutputTokens int64    `json:"output_tokens"`
	TotalTokens  int64    `json:"total_tokens"`
	InputCost    *float64 `json:"input_cost" extensions:"x-nullable"`
	OutputCost   *float64 `json:"output_cost" extensions:"x-nullable"`
	TotalCost    *float64 `json:"total_cost" extensions:"x-nullable"`
}

LabelUsage holds per-label token usage aggregates. A request carrying several labels contributes its full totals to each of them, so label rows overlap and do not sum to the period totals.

type LiveEventPublisher 

type LiveEventPublisher interface {
	PublishUsageEvent(eventType string, entry *UsageEntry)
}

LiveEventPublisher receives compact usage snapshots for realtime dashboard preview. Implementations must not block request handling.

type Logger 

type Logger struct {
	// contains filtered or unexported fields
}

Logger provides async buffered logging with batch writes. It collects usage entries in a channel and flushes them to storage either when the buffer is full or at regular intervals.

func NewLogger 

func NewLogger(store UsageStore, cfg Config) *Logger

NewLogger creates a new async buffered Logger. The logger starts a background goroutine for flushing entries.

func (*Logger) Close 

func (l *Logger) Close() error

Close stops the logger and flushes remaining entries. This should be called during graceful shutdown. Close is idempotent - calling it multiple times is safe.

func (*Logger) Config 

func (l *Logger) Config() Config

Config returns the logger configuration

func (*Logger) SetLivePublisher 

func (l *Logger) SetLivePublisher(p LiveEventPublisher)

SetLivePublisher attaches the optional realtime dashboard publisher.

func (*Logger) Write 

func (l *Logger) Write(entry *UsageEntry)

Write queues a usage entry for async writing. This method is non-blocking. If the buffer is full or the logger is closed, the entry is dropped and a warning is logged.

type LoggerInterface 

type LoggerInterface interface {
	Write(entry *UsageEntry)
	Config() Config
	Close() error
}

LoggerInterface defines the interface for loggers (both real and noop)

type ModelUsage 

type ModelUsage struct {
	Model        string   `json:"model"`
	Provider     string   `json:"provider"`
	ProviderName string   `json:"provider_name,omitempty"`
	InputTokens  int64    `json:"input_tokens"`
	OutputTokens int64    `json:"output_tokens"`
	InputCost    *float64 `json:"input_cost"`
	OutputCost   *float64 `json:"output_cost"`
	TotalCost    *float64 `json:"total_cost"`
}

ModelUsage holds per-model token usage aggregates.

type MongoDBReader 

type MongoDBReader struct {
	// contains filtered or unexported fields
}

MongoDBReader implements UsageReader for MongoDB.

func NewMongoDBReader 

func NewMongoDBReader(database *mongo.Database) (*MongoDBReader, error)

func (*MongoDBReader) GetCacheOverview 

func (r *MongoDBReader) GetCacheOverview(ctx context.Context, params UsageQueryParams) (*CacheOverview, error)

GetCacheOverview returns cached-only aggregates for the admin dashboard.

func (*MongoDBReader) GetDailyUsage 

func (r *MongoDBReader) GetDailyUsage(ctx context.Context, params UsageQueryParams) ([]DailyUsage, error)

GetDailyUsage returns usage statistics grouped by time period (daily, weekly, monthly, yearly).

func (*MongoDBReader) GetSummary 

func (r *MongoDBReader) GetSummary(ctx context.Context, params UsageQueryParams) (*UsageSummary, error)

GetSummary returns aggregated usage statistics for the given query parameters.

func (*MongoDBReader) GetTokenThroughput 

func (r *MongoDBReader) GetTokenThroughput(ctx context.Context, gran ThroughputGranularity, end time.Time, offset int64) (*TokenThroughput, error)

GetTokenThroughput returns the trailing window of token-volume buckets for the overview live-throughput chart. Documents are streamed within the window and bucketed in Go (the prompt-cache split lives in raw_data), mirroring the summary input-segment pass. See foldThroughput in throughput.go for why this streams-and-folds rather than grouping in the database (and TODO(perf) there).

func (*MongoDBReader) GetUsageByLabel 

func (r *MongoDBReader) GetUsageByLabel(ctx context.Context, params UsageQueryParams) ([]LabelUsage, error)

GetUsageByLabel returns token and cost totals grouped by request label. $unwind expands each document's labels array, so a document with several labels contributes its totals to each of them; documents without labels are dropped by $unwind.

func (*MongoDBReader) GetUsageByModel 

func (r *MongoDBReader) GetUsageByModel(ctx context.Context, params UsageQueryParams) ([]ModelUsage, error)

GetUsageByModel returns token and cost totals grouped by model and provider.

func (*MongoDBReader) GetUsageByRequestIDs 

func (r *MongoDBReader) GetUsageByRequestIDs(ctx context.Context, requestIDs []string) (map[string][]UsageLogEntry, error)

GetUsageByRequestIDs returns usage entries grouped by request ID.

func (*MongoDBReader) GetUsageByUserPath 

func (r *MongoDBReader) GetUsageByUserPath(ctx context.Context, params UsageQueryParams) ([]UserPathUsage, error)

GetUsageByUserPath returns token and cost totals grouped by tracked user path.

func (*MongoDBReader) GetUsageLog 

func (r *MongoDBReader) GetUsageLog(ctx context.Context, params UsageLogParams) (*UsageLogResult, error)

GetUsageLog returns a paginated list of individual usage log entries.

type MongoDBStore 

type MongoDBStore struct {
	// contains filtered or unexported fields
}

MongoDBStore implements UsageStore for MongoDB.

func NewMongoDBStore 

func NewMongoDBStore(database *mongo.Database, retentionDays int) (*MongoDBStore, error)

NewMongoDBStore creates a new MongoDB usage store. It creates the collection and indexes if they don't exist. MongoDB handles TTL-based cleanup automatically via TTL indexes.

func (*MongoDBStore) Close 

func (s *MongoDBStore) Close() error

Close is a no-op for MongoDB as the client is managed by the storage layer.

func (*MongoDBStore) Flush 

func (s *MongoDBStore) Flush(_ context.Context) error

Flush is a no-op for MongoDB as writes are synchronous.

func (*MongoDBStore) RecalculatePricing 

func (s *MongoDBStore) RecalculatePricing(ctx context.Context, params RecalculatePricingParams, resolver PricingResolver) (RecalculatePricingResult, error)

RecalculatePricing updates matching MongoDB usage documents with costs computed from the supplied pricing resolver.

func (*MongoDBStore) WriteBatch 

func (s *MongoDBStore) WriteBatch(ctx context.Context, entries []*UsageEntry) error

WriteBatch writes multiple usage entries to MongoDB using InsertMany.

type NoopLogger 

type NoopLogger struct {
	// contains filtered or unexported fields
}

NoopLogger is a logger that does nothing (used when usage tracking is disabled)

func NewNoopLogger 

func NewNoopLogger(cfg Config) *NoopLogger

NewNoopLogger creates a disabled logger that still carries policy config such as whether streaming requests should ask providers to include usage.

func (*NoopLogger) Close 

func (l *NoopLogger) Close() error

Close does nothing

func (*NoopLogger) Config 

func (l *NoopLogger) Config() Config

Config returns the effective config with logging disabled.

func (*NoopLogger) Write 

func (l *NoopLogger) Write(_ *UsageEntry)

Write does nothing

type PartialWriteError 

type PartialWriteError struct {
	TotalEntries int
	FailedCount  int
	Cause        mongo.BulkWriteException
}

PartialWriteError wraps a mongo.BulkWriteException with additional context about how many entries failed vs succeeded.

func (*PartialWriteError) Error 

func (e *PartialWriteError) Error() string

func (*PartialWriteError) Unwrap 

func (e *PartialWriteError) Unwrap() error

type PostgreSQLReader 

type PostgreSQLReader struct {
	// contains filtered or unexported fields
}

PostgreSQLReader implements UsageReader for PostgreSQL databases.

func NewPostgreSQLReader 

func NewPostgreSQLReader(pool *pgxpool.Pool) (*PostgreSQLReader, error)

NewPostgreSQLReader creates a new PostgreSQL usage reader.

func (*PostgreSQLReader) GetCacheOverview 

func (r *PostgreSQLReader) GetCacheOverview(ctx context.Context, params UsageQueryParams) (*CacheOverview, error)

GetCacheOverview returns cached-only aggregates for the admin dashboard.

func (*PostgreSQLReader) GetDailyUsage 

func (r *PostgreSQLReader) GetDailyUsage(ctx context.Context, params UsageQueryParams) ([]DailyUsage, error)

GetDailyUsage returns usage statistics grouped by time period (daily, weekly, monthly, yearly).

func (*PostgreSQLReader) GetSummary 

func (r *PostgreSQLReader) GetSummary(ctx context.Context, params UsageQueryParams) (*UsageSummary, error)

GetSummary returns aggregated usage statistics for the given query parameters.

func (*PostgreSQLReader) GetTokenThroughput 

func (r *PostgreSQLReader) GetTokenThroughput(ctx context.Context, gran ThroughputGranularity, end time.Time, offset int64) (*TokenThroughput, error)

GetTokenThroughput returns the trailing window of token-volume buckets for the overview live-throughput chart. Buckets are epoch-aligned; the prompt-cache split is folded in Go (it lives in raw_data), mirroring the summary path.

func (*PostgreSQLReader) GetUsageByLabel 

func (r *PostgreSQLReader) GetUsageByLabel(ctx context.Context, params UsageQueryParams) ([]LabelUsage, error)

GetUsageByLabel returns token and cost totals grouped by request label. jsonb_array_elements_text expands each row's labels JSONB array, so a row with several labels contributes its totals to each of them; rows with NULL or non-array labels are omitted by the jsonb_typeof guard.

func (*PostgreSQLReader) GetUsageByModel 

func (r *PostgreSQLReader) GetUsageByModel(ctx context.Context, params UsageQueryParams) ([]ModelUsage, error)

GetUsageByModel returns token and cost totals grouped by model and provider.

func (*PostgreSQLReader) GetUsageByRequestIDs 

func (r *PostgreSQLReader) GetUsageByRequestIDs(ctx context.Context, requestIDs []string) (map[string][]UsageLogEntry, error)

GetUsageByRequestIDs returns usage entries grouped by request ID.

func (*PostgreSQLReader) GetUsageByUserPath 

func (r *PostgreSQLReader) GetUsageByUserPath(ctx context.Context, params UsageQueryParams) ([]UserPathUsage, error)

GetUsageByUserPath returns token and cost totals grouped by tracked user path.

func (*PostgreSQLReader) GetUsageLog 

func (r *PostgreSQLReader) GetUsageLog(ctx context.Context, params UsageLogParams) (*UsageLogResult, error)

GetUsageLog returns a paginated list of individual usage log entries.

type PostgreSQLStore 

type PostgreSQLStore struct {
	// contains filtered or unexported fields
}

PostgreSQLStore implements UsageStore for PostgreSQL databases.

func NewPostgreSQLStore 

func NewPostgreSQLStore(pool *pgxpool.Pool, retentionDays int) (*PostgreSQLStore, error)

NewPostgreSQLStore creates a new PostgreSQL usage store. It creates the usage table if it doesn't exist and starts a background cleanup goroutine if retention is configured.

func (*PostgreSQLStore) Close 

func (s *PostgreSQLStore) Close() error

Close stops the cleanup goroutine. Note: We don't close the pool here as it's managed by the storage layer. Safe to call multiple times.

func (*PostgreSQLStore) Flush 

func (s *PostgreSQLStore) Flush(_ context.Context) error

Flush is a no-op for PostgreSQL as writes are synchronous.

func (*PostgreSQLStore) RecalculatePricing 

func (s *PostgreSQLStore) RecalculatePricing(ctx context.Context, params RecalculatePricingParams, resolver PricingResolver) (RecalculatePricingResult, error)

RecalculatePricing updates matching PostgreSQL usage rows with costs computed from the supplied pricing resolver.

func (*PostgreSQLStore) WriteBatch 

func (s *PostgreSQLStore) WriteBatch(ctx context.Context, entries []*UsageEntry) error

WriteBatch writes multiple usage entries to PostgreSQL using batch insert.

type PricingRecalculator 

type PricingRecalculator interface {
	RecalculatePricing(ctx context.Context, params RecalculatePricingParams, resolver PricingResolver) (RecalculatePricingResult, error)
}

PricingRecalculator updates persisted usage cost fields from current pricing metadata.

func NewPricingRecalculator 

func NewPricingRecalculator(store storage.Storage) (PricingRecalculator, error)

NewPricingRecalculator creates a PricingRecalculator from a storage backend. Returns nil if storage is nil.

type PricingResolver 

type PricingResolver interface {
	ResolvePricing(model, providerType string) *core.ModelPricing
}

PricingResolver resolves pricing metadata for a given model and provider type. Implementations should check the registry first and fall back to a reverse-index lookup when the model ID in the usage DB differs from the registry key.

type RecalculatePricingParams 

type RecalculatePricingParams struct {
	UsageQueryParams
}

RecalculatePricingParams identifies the stored usage rows whose costs should be recalculated from the latest model pricing metadata. Row selection (date range, model, provider, label, user path) rides on the embedded UsageQueryParams, sharing the readers' filter semantics.

type RecalculatePricingResult 

type RecalculatePricingResult struct {
	Status         string `json:"status"`
	Matched        int64  `json:"matched"`
	Recalculated   int64  `json:"recalculated"`
	WithPricing    int64  `json:"with_pricing"`
	WithoutPricing int64  `json:"without_pricing"`
}

RecalculatePricingResult summarizes a pricing recalculation run.

type RequestUsageLoader 

type RequestUsageLoader interface {
	GetUsageByRequestIDs(ctx context.Context, requestIDs []string) (map[string][]UsageLogEntry, error)
}

RequestUsageLoader is the slice of the usage reader needed to summarize usage per request.

type RequestUsageSummary 

type RequestUsageSummary struct {
	Entries                   int      `json:"entries"`
	InputTokens               int64    `json:"input_tokens"`
	UncachedInputTokens       int64    `json:"uncached_input_tokens"`
	CachedInputTokens         int64    `json:"cached_input_tokens"`
	CacheWriteInputTokens     int64    `json:"cache_write_input_tokens"`
	OutputTokens              int64    `json:"output_tokens"`
	TotalTokens               int64    `json:"total_tokens"`
	CachedInputRatio          float64  `json:"cached_input_ratio"`
	EstimatedCachedCharacters int64    `json:"estimated_cached_characters"`
	RewriteTokensSaved        int64    `json:"rewrite_tokens_saved,omitempty"`
	RewriteCostSaved          *float64 `json:"rewrite_cost_saved,omitempty"`
}

RequestUsageSummary aggregates usage records that belong to one request ID. InputTokens and TotalTokens are normalized prompt/total counts across providers: cached prompt reads and cache writes are included even when the upstream provider reports them outside the base input token count.

func SummarizeRequestUsage 

func SummarizeRequestUsage(entries []UsageLogEntry) *RequestUsageSummary

SummarizeRequestUsage aggregates one request's usage entries into a normalized summary.

type Result 

type Result struct {
	Logger  LoggerInterface
	Storage storage.Storage
}

Result holds the initialized usage logger and its dependencies. The caller is responsible for calling Close() to release resources.

func New 

func New(ctx context.Context, cfg *config.Config) (*Result, error)

New creates a usage logger from configuration. Returns a Result containing the logger and storage for lifecycle management. The caller must call Result.Close() during shutdown.

If usage tracking is disabled in the config, returns a NoopLogger with nil storage.

func NewWithSharedStorage 

func NewWithSharedStorage(ctx context.Context, cfg *config.Config, store storage.Storage) (*Result, error)

NewWithSharedStorage creates a usage logger using a shared storage connection. This is useful when you want to share the database connection with audit logging. The caller is responsible for closing the storage separately.

func (*Result) Close 

func (r *Result) Close() error

Close releases all resources held by the usage logger. Safe to call multiple times.

type SQLiteReader 

type SQLiteReader struct {
	// contains filtered or unexported fields
}

SQLiteReader implements UsageReader for SQLite databases.

func NewSQLiteReader 

func NewSQLiteReader(db *sql.DB) (*SQLiteReader, error)

NewSQLiteReader creates a new SQLite usage reader.

func (*SQLiteReader) GetCacheOverview 

func (r *SQLiteReader) GetCacheOverview(ctx context.Context, params UsageQueryParams) (*CacheOverview, error)

GetCacheOverview returns cached-only aggregates for the admin dashboard.

func (*SQLiteReader) GetDailyUsage 

func (r *SQLiteReader) GetDailyUsage(ctx context.Context, params UsageQueryParams) ([]DailyUsage, error)

GetDailyUsage returns usage statistics grouped by time period (daily, weekly, monthly, yearly).

func (*SQLiteReader) GetSummary 

func (r *SQLiteReader) GetSummary(ctx context.Context, params UsageQueryParams) (*UsageSummary, error)

GetSummary returns aggregated usage statistics for the given query parameters.

func (*SQLiteReader) GetTokenThroughput 

func (r *SQLiteReader) GetTokenThroughput(ctx context.Context, gran ThroughputGranularity, end time.Time, offset int64) (*TokenThroughput, error)

GetTokenThroughput returns the trailing window of token-volume buckets for the overview live-throughput chart. Buckets are epoch-aligned; the prompt-cache split is folded in Go (it lives in raw_data), mirroring the summary path.

func (*SQLiteReader) GetUsageByLabel 

func (r *SQLiteReader) GetUsageByLabel(ctx context.Context, params UsageQueryParams) ([]LabelUsage, error)

GetUsageByLabel returns token and cost totals grouped by request label. json_each expands each row's labels JSON array, so a row with several labels contributes its totals to each of them; unlabelled rows are omitted.

func (*SQLiteReader) GetUsageByModel 

func (r *SQLiteReader) GetUsageByModel(ctx context.Context, params UsageQueryParams) ([]ModelUsage, error)

GetUsageByModel returns token and cost totals grouped by model and provider.

func (*SQLiteReader) GetUsageByRequestIDs 

func (r *SQLiteReader) GetUsageByRequestIDs(ctx context.Context, requestIDs []string) (map[string][]UsageLogEntry, error)

GetUsageByRequestIDs returns usage entries grouped by request ID.

func (*SQLiteReader) GetUsageByUserPath 

func (r *SQLiteReader) GetUsageByUserPath(ctx context.Context, params UsageQueryParams) ([]UserPathUsage, error)

GetUsageByUserPath returns token and cost totals grouped by tracked user path.

func (*SQLiteReader) GetUsageLog 

func (r *SQLiteReader) GetUsageLog(ctx context.Context, params UsageLogParams) (*UsageLogResult, error)

GetUsageLog returns a paginated list of individual usage log entries.

type SQLiteStore 

type SQLiteStore struct {
	// contains filtered or unexported fields
}

SQLiteStore implements UsageStore for SQLite databases.

func NewSQLiteStore 

func NewSQLiteStore(db *sql.DB, retentionDays int) (*SQLiteStore, error)

NewSQLiteStore creates a new SQLite usage store. It creates the usage table if it doesn't exist and starts a background cleanup goroutine if retention is configured.

func (*SQLiteStore) Close 

func (s *SQLiteStore) Close() error

Close stops the cleanup goroutine. Note: We don't close the DB here as it's managed by the storage layer. Safe to call multiple times.

func (*SQLiteStore) Flush 

func (s *SQLiteStore) Flush(_ context.Context) error

Flush is a no-op for SQLite as writes are synchronous.

func (*SQLiteStore) RecalculatePricing 

func (s *SQLiteStore) RecalculatePricing(ctx context.Context, params RecalculatePricingParams, resolver PricingResolver) (RecalculatePricingResult, error)

RecalculatePricing updates matching SQLite usage rows with costs computed from the supplied pricing resolver.

func (*SQLiteStore) WriteBatch 

func (s *SQLiteStore) WriteBatch(ctx context.Context, entries []*UsageEntry) error

WriteBatch writes multiple usage entries to SQLite using batch insert. Entries are chunked to stay within SQLite's parameter limit.

type StreamUsageObserver 

type StreamUsageObserver struct {
	// contains filtered or unexported fields
}

StreamUsageObserver extracts usage data from parsed SSE JSON payloads.

func NewStreamUsageObserver 

func NewStreamUsageObserver(logger LoggerInterface, model, provider, requestID, endpoint string, pricingResolver PricingResolver, userPath ...string) *StreamUsageObserver

func (*StreamUsageObserver) OnJSONEvent 

func (o *StreamUsageObserver) OnJSONEvent(chunk map[string]any)

func (*StreamUsageObserver) OnStreamClose 

func (o *StreamUsageObserver) OnStreamClose()

func (*StreamUsageObserver) SetLabels 

func (o *StreamUsageObserver) SetLabels(labels []string)

SetLabels attaches the request labels extracted from tagging headers.

func (*StreamUsageObserver) SetProviderName 

func (o *StreamUsageObserver) SetProviderName(providerName string)

func (*StreamUsageObserver) SetRewriteTokensSaved 

func (o *StreamUsageObserver) SetRewriteTokensSaved(tokensSaved int)

SetRewriteTokensSaved attaches the request's rewrite savings estimate so the usage entry extracted from the stream records it (with its cost, when pricing is resolvable).

func (*StreamUsageObserver) WantsJSONEvent 

func (o *StreamUsageObserver) WantsJSONEvent(raw []byte) bool

WantsJSONEvent reports whether the raw SSE payload can carry usage data. This lets the observed stream skip JSON decoding for the vast majority of content-delta chunks. Known trade-off: a provider that JSON-escapes key characters (e.g. "usage") would slip past this byte scan; no known provider does, and covering it would mean decoding every chunk again.

type ThroughputBucket 

type ThroughputBucket struct {
	Start               time.Time `json:"start"`
	InputTokens         int64     `json:"input_tokens"`
	OutputTokens        int64     `json:"output_tokens"`
	PromptCachedTokens  int64     `json:"prompt_cached_tokens"`
	LocallyCachedTokens int64     `json:"locally_cached_tokens"`
}

ThroughputBucket is one time bucket of token volume, split into the four series the overview live chart stacks. Start is the bucket's inclusive start (UTC); the bucket spans the parent TokenThroughput's BucketSeconds.

type ThroughputGranularity 

type ThroughputGranularity struct {
	Name        string
	BucketSize  time.Duration
	WindowCount int
}

ThroughputGranularity describes one bucket width for the overview live token-throughput chart. WindowCount is how many trailing buckets the chart shows, so the window spans BucketSize * WindowCount.

func ParseThroughputGranularity 

func ParseThroughputGranularity(name string) (ThroughputGranularity, error)

ParseThroughputGranularity resolves a public granularity name (second, minute, hour, day) to its bucket configuration.

type TokenThroughput 

type TokenThroughput struct {
	Granularity   string             `json:"granularity"`
	BucketSeconds int                `json:"bucket_seconds"`
	Buckets       []ThroughputBucket `json:"buckets"`
}

TokenThroughput is a fixed-width window of token-volume buckets ending at the current time, powering the overview live-throughput chart.

func EmptyTokenThroughput 

func EmptyTokenThroughput(gran ThroughputGranularity, end time.Time, offset int64) *TokenThroughput

EmptyTokenThroughput returns a zero-filled window, used when usage tracking is disabled so the dashboard still renders an (empty) chart.

type UsageEntry 

type UsageEntry struct {
	// ID is a unique identifier for this usage entry (UUID)
	ID string `json:"id" bson:"_id"`

	// RequestID links to the audit log entry (from X-Request-ID header)
	RequestID string `json:"request_id" bson:"request_id"`

	// ProviderID is the provider's response ID (e.g., "chatcmpl-abc123", "msg_xyz")
	ProviderID string `json:"provider_id" bson:"provider_id"`

	// Timestamp is when the request completed
	Timestamp time.Time `json:"timestamp" bson:"timestamp"`

	// Request context
	Model        string `json:"model" bson:"model"`
	Provider     string `json:"provider" bson:"provider"` // canonical provider type used for routing, filters, and pricing
	ProviderName string `json:"provider_name,omitempty" bson:"provider_name,omitempty"`
	Endpoint     string `json:"endpoint" bson:"endpoint"`
	UserPath     string `json:"user_path,omitempty" bson:"user_path,omitempty"`
	CacheType    string `json:"cache_type,omitempty" bson:"cache_type,omitempty"`

	// Labels are request labels extracted from configured tagging headers.
	Labels []string `json:"labels,omitempty" bson:"labels,omitempty"`

	// Standard token counts (normalized across providers)
	InputTokens  int `json:"input_tokens" bson:"input_tokens"`
	OutputTokens int `json:"output_tokens" bson:"output_tokens"`
	TotalTokens  int `json:"total_tokens" bson:"total_tokens"`

	// RawData contains provider-specific extended usage data (JSONB)
	// Examples:
	//   OpenAI: {"cached_tokens": 100, "reasoning_tokens": 50}
	//   Anthropic: {"cache_creation_input_tokens": 200, "cache_read_input_tokens": 150}
	//   Gemini: {"cached_tokens": 100, "thought_tokens": 75, "tool_use_tokens": 25}
	RawData map[string]any `json:"raw_data,omitempty" bson:"raw_data,omitempty"`

	// Cost fields (nil = unknown/model not in list, 0.0 = free)
	InputCost  *float64 `json:"input_cost,omitempty" bson:"input_cost,omitempty"`
	OutputCost *float64 `json:"output_cost,omitempty" bson:"output_cost,omitempty"`
	TotalCost  *float64 `json:"total_cost,omitempty" bson:"total_cost,omitempty"`
	CostSource string   `json:"cost_source,omitempty" bson:"cost_source,omitempty"`

	// CostsCalculationCaveat describes any incomplete aspects of cost calculation.
	// Empty means all token types were fully mapped to pricing data.
	CostsCalculationCaveat string `json:"costs_calculation_caveat,omitempty" bson:"costs_calculation_caveat,omitempty"`

	// RewriteTokensSaved is the prompt tokens that request rewriters (ext
	// extensions, e.g. token compression) estimate they removed from this
	// request before it reached the provider. Zero when no rewriter reported
	// savings.
	RewriteTokensSaved int `json:"rewrite_tokens_saved,omitempty" bson:"rewrite_tokens_saved,omitempty"`

	// RewriteCostSaved is the estimated input cost those removed tokens
	// would have added, priced with the same model pricing as the request
	// (nil = pricing unknown).
	RewriteCostSaved *float64 `json:"rewrite_cost_saved,omitempty" bson:"rewrite_cost_saved,omitempty"`
}

UsageEntry represents a single token usage record.

func ExtractFromCachedResponseBody 

func ExtractFromCachedResponseBody(
	body []byte,
	requestID, model, provider, endpoint, cacheType string,
	pricing ...*core.ModelPricing,
) *UsageEntry

ExtractFromCachedResponseBody converts a cached OpenAI-compatible response body into a synthetic usage entry for a cache hit. If the response body cannot be parsed, it still returns a minimal zero-token entry so cache-hit counts remain observable.

func ExtractFromChatResponse 

func ExtractFromChatResponse(resp *core.ChatResponse, requestID, provider, endpoint string, pricing ...*core.ModelPricing) *UsageEntry

ExtractFromChatResponse extracts usage data from a ChatResponse. It normalizes the usage data into a UsageEntry and preserves raw extended data. If pricing is provided, granular cost fields are calculated. For `/v1/batches` endpoints (exact or subpath), batch pricing overrides (BatchInputPerMtok/BatchOutputPerMtok) may replace standard input/output rates.

func ExtractFromEmbeddingResponse 

func ExtractFromEmbeddingResponse(resp *core.EmbeddingResponse, requestID, provider, endpoint string, pricing ...*core.ModelPricing) *UsageEntry

ExtractFromEmbeddingResponse extracts usage data from an EmbeddingResponse. Embeddings only have prompt tokens (no output tokens). For `/v1/batches` endpoints (exact or subpath), BatchInputPerMtok may replace standard InputPerMtok when pricingForEndpoint applies batch overrides.

func ExtractFromRealtimeResponseDone 

func ExtractFromRealtimeResponseDone(payload []byte, requestID, model, provider string, pricing ...*core.ModelPricing) *UsageEntry

ExtractFromRealtimeResponseDone builds a usage entry from a realtime "response.done" event. A realtime session produces one such event per model response, each carrying its own usage, so the caller writes one entry per event. It returns nil when the payload is not a response.done event or carries no usage, so non-billable events (audio deltas, transcripts) are skipped cheaply.

func ExtractFromResponsesResponse 

func ExtractFromResponsesResponse(resp *core.ResponsesResponse, requestID, provider, endpoint string, pricing ...*core.ModelPricing) *UsageEntry

ExtractFromResponsesResponse extracts usage data from a ResponsesResponse. It normalizes the usage data into a UsageEntry and preserves raw extended data. If pricing is provided, cost fields are calculated. For `/v1/batches` endpoints (exact or subpath), batch pricing overrides (BatchInputPerMtok/BatchOutputPerMtok) may replace standard input/output rates.

func ExtractFromSSEUsage 

func ExtractFromSSEUsage(
	providerID string,
	inputTokens, outputTokens, totalTokens int,
	rawData map[string]any,
	requestID, model, provider, endpoint string,
	pricing ...*core.ModelPricing,
) *UsageEntry

ExtractFromSSEUsage creates a UsageEntry from SSE-extracted usage data. This is used for streaming responses where usage is extracted from the final SSE event. If pricing is provided, cost fields are calculated. For `/v1/batches` endpoints (exact or subpath), batch pricing overrides (BatchInputPerMtok/BatchOutputPerMtok) may replace standard input/output rates.

func ExtractFromSpeechRequest 

func ExtractFromSpeechRequest(input string, output []byte, format, requestID, model, provider string, pricing ...*core.ModelPricing) *UsageEntry

ExtractFromSpeechRequest builds a usage entry for a text-to-speech request. Speech responses are binary audio with no provider-reported usage, so the billable units are derived locally: the input character count (per-character models such as tts-1) and the synthesized audio duration (per-second-output models such as gpt-4o-mini-tts), both recorded in RawData so the interaction stays observable and pricing can apply. output is the returned audio and format its response_format/MIME type; duration is measured for wav/pcm/mp3 (see measureSpeechDurationSeconds). model is the resolved route model (not the raw user input) so the row groups and prices consistently with the pricing lookup, mirroring the transcription extractor.

func ExtractFromTranscriptionResponse 

func ExtractFromTranscriptionResponse(body []byte, requestID, model, provider string, pricing ...*core.ModelPricing) *UsageEntry

ExtractFromTranscriptionResponse builds a usage entry for a speech-to-text request. The response body is proxied verbatim; when it is JSON it may carry a usage object (token- or duration-based). The entry is always returned so the interaction stays observable even when the provider reports no usage (whisper, or non-JSON response formats such as text/srt/vtt).

type UsageLogEntry 

type UsageLogEntry struct {
	ID                     string         `json:"id"`
	RequestID              string         `json:"request_id"`
	ProviderID             string         `json:"provider_id"`
	Timestamp              time.Time      `json:"timestamp"`
	Model                  string         `json:"model"`
	Provider               string         `json:"provider"`
	ProviderName           string         `json:"provider_name,omitempty"`
	Endpoint               string         `json:"endpoint"`
	UserPath               string         `json:"user_path,omitempty"`
	CacheType              string         `json:"cache_type,omitempty"`
	Labels                 []string       `json:"labels,omitempty"`
	InputTokens            int            `json:"input_tokens"`
	OutputTokens           int            `json:"output_tokens"`
	TotalTokens            int            `json:"total_tokens"`
	UncachedInputTokens    int64          `json:"uncached_input_tokens,omitempty"`
	CachedInputTokens      int64          `json:"cached_input_tokens,omitempty"`
	CacheWriteInputTokens  int64          `json:"cache_write_input_tokens,omitempty"`
	CachedInputRatio       float64        `json:"cached_input_ratio,omitempty"`
	InputCost              *float64       `json:"input_cost"`
	OutputCost             *float64       `json:"output_cost"`
	TotalCost              *float64       `json:"total_cost"`
	CostSource             string         `json:"cost_source,omitempty"`
	RawData                map[string]any `json:"raw_data,omitempty"`
	CostsCalculationCaveat string         `json:"costs_calculation_caveat,omitempty"`
	RewriteTokensSaved     int64          `json:"rewrite_tokens_saved,omitempty"`
	RewriteCostSaved       *float64       `json:"rewrite_cost_saved,omitempty"`
}

UsageLogEntry represents a single usage record in the request log.

The cached_input_tokens, uncached_input_tokens, cache_write_input_tokens, and cached_input_ratio fields are derived from RawData at read time by EnrichUsageLogEntry; storage layers never populate them.

type UsageLogParams 

type UsageLogParams struct {
	UsageQueryParams        // embed date range and data filters
	Search           string // free-text search on model/provider/request_id
	Limit            int    // page size (default 50, max 200)
	Offset           int    // pagination offset
}

UsageLogParams specifies query parameters for paginated usage log retrieval. Data filters (model, provider, label, user path) live on the embedded UsageQueryParams; only the log-specific view options are declared here.

type UsageLogResult 

type UsageLogResult struct {
	Entries []UsageLogEntry `json:"entries"`
	Total   int             `json:"total"`
	Limit   int             `json:"limit"`
	Offset  int             `json:"offset"`
}

UsageLogResult holds a paginated list of usage log entries.

type UsageQueryParams 

type UsageQueryParams struct {
	StartDate time.Time // Inclusive start (day precision)
	EndDate   time.Time // Inclusive end (day precision)
	Interval  string    // "daily", "weekly", "monthly", "yearly"
	TimeZone  string    // IANA timezone used for day-boundary interpretation and grouping
	UserPath  string    // subtree filter on tracked user path
	Model     string    // filter by exact model name (optional)
	Provider  string    // filter by provider name or provider type (optional)
	Label     string    // filter by request label, exact match (optional)
	CacheMode string    // "uncached" (default), "cached", or "all"
}

UsageQueryParams specifies the query parameters for usage data retrieval. The optional filters (UserPath, Model, Provider, Label) apply uniformly to every reader method, so summaries, breakdowns, and the request log all describe the same filtered slice of traffic.

type UsageReader 

type UsageReader interface {
	// GetSummary returns aggregated usage statistics for the given date range.
	// If both StartDate and EndDate are zero, returns all-time statistics.
	GetSummary(ctx context.Context, params UsageQueryParams) (*UsageSummary, error)

	// GetDailyUsage returns usage statistics grouped by the specified interval.
	// If both StartDate and EndDate are zero, returns all available data.
	GetDailyUsage(ctx context.Context, params UsageQueryParams) ([]DailyUsage, error)

	// GetUsageByModel returns per-model token usage aggregates for the given date range.
	GetUsageByModel(ctx context.Context, params UsageQueryParams) ([]ModelUsage, error)

	// GetUsageByUserPath returns per-user-path token usage aggregates for the given date range.
	GetUsageByUserPath(ctx context.Context, params UsageQueryParams) ([]UserPathUsage, error)

	// GetUsageByLabel returns per-label token usage aggregates for the given
	// date range. Unlabelled entries are omitted; entries with several labels
	// count once per label.
	GetUsageByLabel(ctx context.Context, params UsageQueryParams) ([]LabelUsage, error)

	// GetUsageLog returns a paginated list of individual usage entries with optional filtering.
	GetUsageLog(ctx context.Context, params UsageLogParams) (*UsageLogResult, error)

	// GetUsageByRequestIDs returns usage log entries grouped by request_id.
	// Missing IDs are omitted from the returned map.
	GetUsageByRequestIDs(ctx context.Context, requestIDs []string) (map[string][]UsageLogEntry, error)

	// GetCacheOverview returns cached-only aggregates for the admin dashboard.
	GetCacheOverview(ctx context.Context, params UsageQueryParams) (*CacheOverview, error)

	// GetTokenThroughput returns a fixed-width window of token-volume buckets
	// (input/output/prompt-cached/locally-cached) ending at end, for the
	// overview live-throughput chart. The window is global (not user-path
	// scoped). offset is the request timezone's offset from UTC in seconds, so
	// buckets align to local boundaries (e.g. day buckets at local midnight).
	GetTokenThroughput(ctx context.Context, gran ThroughputGranularity, end time.Time, offset int64) (*TokenThroughput, error)
}

UsageReader provides read access to usage data for the admin API.

func NewReader 

func NewReader(store storage.Storage) (UsageReader, error)

NewReader creates a UsageReader from a storage backend. Returns nil if the storage is nil (usage data not available).

type UsageStore 

type UsageStore interface {
	// WriteBatch writes multiple usage entries to storage.
	// This is called by the Logger when flushing buffered entries.
	WriteBatch(ctx context.Context, entries []*UsageEntry) error

	// Flush forces any pending writes to complete.
	// Called during graceful shutdown.
	Flush(ctx context.Context) error

	// Close releases resources and flushes pending writes.
	Close() error
}

UsageStore defines the interface for usage storage backends. Implementations must be safe for concurrent use.

type UsageSummary 

type UsageSummary struct {
	TotalRequests         int      `json:"total_requests"`
	TotalInput            int64    `json:"total_input_tokens"`
	TotalOutput           int64    `json:"total_output_tokens"`
	TotalTokens           int64    `json:"total_tokens"`
	UncachedInputTokens   int64    `json:"uncached_input_tokens"`
	CachedInputTokens     int64    `json:"cached_input_tokens"`
	CacheWriteInputTokens int64    `json:"cache_write_input_tokens"`
	TotalInputCost        *float64 `json:"total_input_cost"`
	TotalOutputCost       *float64 `json:"total_output_cost"`
	TotalCost             *float64 `json:"total_cost"`
	// Rewrite savings: prompt tokens request rewriters removed before the
	// provider call, and the estimated input cost avoided (nil when no
	// matched row had a priced savings estimate).
	RewriteTokensSaved int64    `json:"rewrite_tokens_saved"`
	RewriteCostSaved   *float64 `json:"rewrite_cost_saved"`
}

UsageSummary holds aggregated usage statistics over a time period.

The *_input_tokens split fields (uncached/cached/cache-write) are the provider prompt-cache breakdown of the input, summed per row via addInputSegments which reuses EntryInputSegments. Their sum is the provider-side "input parts" total; for additive-accounting providers (Anthropic) it exceeds TotalInput, which only sums the input_tokens column. Storage layers populate them by streaming rows; they are zero when the reader is disabled.

type UserPathUsage 

type UserPathUsage struct {
	UserPath     string   `json:"user_path"`
	InputTokens  int64    `json:"input_tokens"`
	OutputTokens int64    `json:"output_tokens"`
	TotalTokens  int64    `json:"total_tokens"`
	InputCost    *float64 `json:"input_cost" extensions:"x-nullable"`
	OutputCost   *float64 `json:"output_cost" extensions:"x-nullable"`
	TotalCost    *float64 `json:"total_cost" extensions:"x-nullable"`
}

UserPathUsage holds per-user-path token usage aggregates.

Source Files

View all Source files

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.