internal/gateway/usage.go

Rama main · réplica de lectura del repositorio original.

package gateway

import (
	"context"
	"strings"

	"github.com/enterpilot/gomodel/internal/core"
	"github.com/enterpilot/gomodel/internal/usage"
)

// LogUsage writes one non-streaming usage entry when usage is enabled.
func (o *InferenceOrchestrator) LogUsage(
	ctx context.Context,
	workflow *core.Workflow,
	model, providerType, providerName string,
	extractFn func(*core.ModelPricing) *usage.UsageEntry,
) {
	o.logUsage(ctx, workflow, model, providerType, providerName, extractFn)
}

func (o *InferenceOrchestrator) logUsage(
	ctx context.Context,
	workflow *core.Workflow,
	model, providerType, providerName string,
	extractFn func(*core.ModelPricing) *usage.UsageEntry,
) {
	if o.usageLogger == nil || !o.usageLogger.Config().Enabled || (workflow != nil && !workflow.UsageEnabled()) {
		return
	}
	var pricing *core.ModelPricing
	if o.pricingResolver != nil {
		pricing = o.pricingResolver.ResolvePricing(model, effectivePricingProvider(providerType, providerName))
	}
	if entry := extractFn(pricing); entry != nil {
		entry.ProviderName = strings.TrimSpace(providerName)
		entry.RoutedModel = strings.TrimSpace(model)
		entry.UserPath = core.UserPathFromContext(ctx)
		entry.SessionID = core.SessionIDFromContext(ctx)
		entry.AuthKeyID = core.GetAuthKeyID(ctx)
		entry.Labels = core.RequestLabelsFromContext(ctx)
		usage.ApplyRewriteSavings(entry, core.RewriteTokensSavedFromContext(ctx), pricing)
		o.usageLogger.Write(entry)
	}
}

func effectivePricingProvider(providerType, providerName string) string {
	if name := strings.TrimSpace(providerName); name != "" {
		return name
	}
	return strings.TrimSpace(providerType)
}

// ShouldEnforceReturningUsageData reports whether streams should request usage chunks.
func (o *InferenceOrchestrator) ShouldEnforceReturningUsageData() bool {
	if o.usageLogger == nil {
		return false
	}
	cfg := o.usageLogger.Config()
	return cfg.Enabled && cfg.EnforceReturningUsageData
}