Unify provider/model normalization for AI cost export

This commit is contained in:
rcourtman 2025-12-12 13:04:42 +00:00
parent 5c9ebf8d3a
commit 11ebc0e91c
3 changed files with 79 additions and 41 deletions

View file

@ -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)
}

View file

@ -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

View file

@ -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,