From 11ebc0e91c43606689e142ef8c055e0c496a7b4e Mon Sep 17 00:00:00 2001 From: rcourtman Date: Fri, 12 Dec 2025 13:04:42 +0000 Subject: [PATCH] Unify provider/model normalization for AI cost export --- internal/ai/cost/resolve.go | 53 +++++++++++++++++++++++++++++++++++++ internal/ai/cost/store.go | 43 +++--------------------------- internal/api/ai_handlers.go | 24 ++++++++++++++++- 3 files changed, 79 insertions(+), 41 deletions(-) create mode 100644 internal/ai/cost/resolve.go diff --git a/internal/ai/cost/resolve.go b/internal/ai/cost/resolve.go new file mode 100644 index 0000000..9ec7fcf --- /dev/null +++ b/internal/ai/cost/resolve.go @@ -0,0 +1,53 @@ +package cost + +import "strings" + +// ResolveProviderAndModel normalizes an AI usage event into a provider and model pair for reporting/pricing. +// It intentionally applies compatibility heuristics so OpenAI-compatible APIs (like DeepSeek) still price correctly. +func ResolveProviderAndModel(eventProvider, requestModel, responseModel string) (provider, model string) { + provider = strings.ToLower(strings.TrimSpace(eventProvider)) + + model = normalizeModelForProvider(provider, requestModel, responseModel) + if provider == "" && requestModel != "" { + parts := strings.SplitN(strings.TrimSpace(requestModel), ":", 2) + if len(parts) == 2 && strings.TrimSpace(parts[0]) != "" { + provider = strings.ToLower(strings.TrimSpace(parts[0])) + } + } + + provider, model = inferProviderAndModel(provider, model) + return provider, model +} + +func normalizeModelForProvider(provider, requestModel, responseModel string) string { + if strings.TrimSpace(requestModel) != "" { + parts := strings.SplitN(requestModel, ":", 2) + if len(parts) == 2 && strings.ToLower(strings.TrimSpace(parts[0])) == provider { + return strings.TrimSpace(parts[1]) + } + return strings.TrimSpace(requestModel) + } + if strings.TrimSpace(responseModel) != "" { + parts := strings.SplitN(responseModel, ":", 2) + if len(parts) == 2 && strings.ToLower(strings.TrimSpace(parts[0])) == provider { + return strings.TrimSpace(parts[1]) + } + return strings.TrimSpace(responseModel) + } + return "" +} + +func inferProviderAndModel(provider, model string) (string, string) { + switch provider { + case "openai": + trimmed := strings.ToLower(strings.TrimSpace(model)) + parts := strings.SplitN(trimmed, ":", 2) + if len(parts) == 2 && strings.TrimSpace(parts[0]) == "deepseek" { + return "deepseek", strings.TrimSpace(parts[1]) + } + if strings.HasPrefix(trimmed, "deepseek") { + return "deepseek", strings.TrimSpace(model) + } + } + return provider, strings.TrimSpace(model) +} diff --git a/internal/ai/cost/store.go b/internal/ai/cost/store.go index c1d3217..0fd6e17 100644 --- a/internal/ai/cost/store.go +++ b/internal/ai/cost/store.go @@ -164,9 +164,7 @@ func (s *Store) GetSummary(days int) Summary { var totalInput, totalOutput int64 for _, e := range events { - provider := strings.ToLower(strings.TrimSpace(e.Provider)) - model := normalizeModel(provider, e.RequestModel, e.ResponseModel) - provider, model = inferProviderAndModel(provider, model) + provider, model := ResolveProviderAndModel(e.Provider, e.RequestModel, e.ResponseModel) k := pmKey{provider: provider, model: model} pm := pmTotals[k] @@ -309,37 +307,6 @@ func (s *Store) scheduleSaveLocked() { }) } -func normalizeModel(provider, requestModel, responseModel string) string { - if requestModel != "" { - parts := strings.SplitN(requestModel, ":", 2) - if len(parts) == 2 && parts[0] == provider { - return parts[1] - } - return requestModel - } - if responseModel != "" { - parts := strings.SplitN(responseModel, ":", 2) - if len(parts) == 2 && parts[0] == provider { - return parts[1] - } - return responseModel - } - return "" -} - -func inferProviderAndModel(provider, model string) (string, string) { - if provider == "openai" { - parts := strings.SplitN(strings.TrimSpace(model), ":", 2) - if len(parts) == 2 && strings.ToLower(parts[0]) == "deepseek" { - return "deepseek", parts[1] - } - if strings.HasPrefix(strings.ToLower(strings.TrimSpace(model)), "deepseek") { - return "deepseek", model - } - } - return provider, model -} - // ProviderModelSummary is a rollup for a provider/model pair. type ProviderModelSummary struct { Provider string `json:"provider"` @@ -419,9 +386,7 @@ func summarizeUseCases(events []UsageEvent) []UseCaseSummary { t.input += int64(e.InputTokens) t.output += int64(e.OutputTokens) - provider := strings.ToLower(strings.TrimSpace(e.Provider)) - model := normalizeModel(provider, e.RequestModel, e.ResponseModel) - provider, model = inferProviderAndModel(provider, model) + provider, model := ResolveProviderAndModel(e.Provider, e.RequestModel, e.ResponseModel) usd, known, _ := EstimateUSD(provider, model, int64(e.InputTokens), int64(e.OutputTokens)) if known { @@ -491,9 +456,7 @@ func summarizeTargets(events []UsageEvent) []TargetSummary { t.input += int64(e.InputTokens) t.output += int64(e.OutputTokens) - provider := strings.ToLower(strings.TrimSpace(e.Provider)) - model := normalizeModel(provider, e.RequestModel, e.ResponseModel) - provider, model = inferProviderAndModel(provider, model) + provider, model := ResolveProviderAndModel(e.Provider, e.RequestModel, e.ResponseModel) usd, known, _ := EstimateUSD(provider, model, int64(e.InputTokens), int64(e.OutputTokens)) if known { t.usd += usd diff --git a/internal/api/ai_handlers.go b/internal/api/ai_handlers.go index d4f067e..def16ee 100644 --- a/internal/api/ai_handlers.go +++ b/internal/api/ai_handlers.go @@ -2536,9 +2536,24 @@ func (h *AISettingsHandler) HandleExportAICostHistory(w http.ResponseWriter, r * if format == "json" { w.Header().Set("Content-Type", "application/json") + type exportEvent struct { + cost.UsageEvent + EstimatedUSD float64 `json:"estimated_usd,omitempty"` + PricingKnown bool `json:"pricing_known"` + } + exported := make([]exportEvent, 0, len(events)) + for _, e := range events { + provider, model := cost.ResolveProviderAndModel(e.Provider, e.RequestModel, e.ResponseModel) + usd, ok, _ := cost.EstimateUSD(provider, model, int64(e.InputTokens), int64(e.OutputTokens)) + exported = append(exported, exportEvent{ + UsageEvent: e, + EstimatedUSD: usd, + PricingKnown: ok, + }) + } resp := map[string]any{ "days": days, - "events": events, + "events": exported, } if err := json.NewEncoder(w).Encode(resp); err != nil { log.Error().Err(err).Msg("Failed to write AI cost export JSON") @@ -2556,11 +2571,16 @@ func (h *AISettingsHandler) HandleExportAICostHistory(w http.ResponseWriter, r * "use_case", "input_tokens", "output_tokens", + "estimated_usd", + "pricing_known", "target_type", "target_id", "finding_id", }) for _, e := range events { + provider, model := cost.ResolveProviderAndModel(e.Provider, e.RequestModel, e.ResponseModel) + usd, ok, _ := cost.EstimateUSD(provider, model, int64(e.InputTokens), int64(e.OutputTokens)) + _ = cw.Write([]string{ e.Timestamp.UTC().Format(time.RFC3339Nano), e.Provider, @@ -2569,6 +2589,8 @@ func (h *AISettingsHandler) HandleExportAICostHistory(w http.ResponseWriter, r * e.UseCase, strconv.Itoa(e.InputTokens), strconv.Itoa(e.OutputTokens), + strconv.FormatFloat(usd, 'f', 6, 64), + strconv.FormatBool(ok), e.TargetType, e.TargetID, e.FindingID,