git.iko.soy/crush
on main help sign in
fix: cost calculation when using openrouter
47bde03Kujtim Hoxha
internal/agent/agent.go+49−5
compare with main
unifiedsplit
@@ -17,6 +17,7 @@ import (
17 "charm.land/fantasy/providers/bedrock"
17 "charm.land/fantasy/providers/bedrock"
18 "charm.land/fantasy/providers/google"
18 "charm.land/fantasy/providers/google"
19 "charm.land/fantasy/providers/openai"
19 "charm.land/fantasy/providers/openai"
20+ "charm.land/fantasy/providers/openrouter"
20 "github.com/charmbracelet/catwalk/pkg/catwalk"
21 "github.com/charmbracelet/catwalk/pkg/catwalk"
21 "github.com/charmbracelet/crush/internal/agent/tools"
22 "github.com/charmbracelet/crush/internal/agent/tools"
22 "github.com/charmbracelet/crush/internal/config"
23 "github.com/charmbracelet/crush/internal/config"
@@ -349,7 +350,7 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*fantasy
349 finishReason = message.FinishReasonToolUse
350 finishReason = message.FinishReasonToolUse
350 }
351 }
351 currentAssistant.AddFinish(finishReason, "", "")
352 currentAssistant.AddFinish(finishReason, "", "")
352− a.updateSessionUsage(a.largeModel, &currentSession, stepResult.Usage)
353+ a.updateSessionUsage(a.largeModel, &currentSession, stepResult.Usage, a.openrouterCost(stepResult.ProviderMetadata))
353 sessionLock.Lock()
354 sessionLock.Lock()
354 _, sessionErr := a.sessions.Save(genCtx, currentSession)
355 _, sessionErr := a.sessions.Save(genCtx, currentSession)
355 sessionLock.Unlock()
356 sessionLock.Unlock()
@@ -562,7 +563,19 @@ func (a *sessionAgent) Summarize(ctx context.Context, sessionID string, opts fan
562 return err
563 return err
563 }
564 }
564
565
566+ var openrouterCost *float64
567+ for _, step := range resp.Steps {
568+ stepCost := a.openrouterCost(step.ProviderMetadata)
569+ if stepCost != nil {
570+ newCost := *stepCost
571+ if openrouterCost != nil {
572+ newCost += *openrouterCost
573+ }
574+ openrouterCost = &newCost
575+ }
576+ }
577+
565− a.updateSessionUsage(a.largeModel, &currentSession, resp.TotalUsage)
578+ a.updateSessionUsage(a.largeModel, &currentSession, resp.TotalUsage, openrouterCost)
566
579
567 // just in case get just the last usage
580 // just in case get just the last usage
568 usage := resp.Response.Usage
581 usage := resp.Response.Usage
@@ -690,7 +703,20 @@ func (a *sessionAgent) generateTitle(ctx context.Context, session *session.Sessi
690 }
703 }
691
704
692 session.Title = title
705 session.Title = title
706+
707+ var openrouterCost *float64
708+ for _, step := range resp.Steps {
709+ stepCost := a.openrouterCost(step.ProviderMetadata)
710+ if stepCost != nil {
711+ newCost := *stepCost
712+ if openrouterCost != nil {
713+ newCost += *openrouterCost
714+ }
715+ openrouterCost = &newCost
716+ }
717+ }
718+
693− a.updateSessionUsage(a.smallModel, session, resp.TotalUsage)
719+ a.updateSessionUsage(a.smallModel, session, resp.TotalUsage, openrouterCost)
694 _, saveErr := a.sessions.Save(ctx, *session)
720 _, saveErr := a.sessions.Save(ctx, *session)
695 if saveErr != nil {
721 if saveErr != nil {
696 slog.Error("failed to save session title & usage", "error", saveErr)
722 slog.Error("failed to save session title & usage", "error", saveErr)
@@ -698,7 +724,20 @@ func (a *sessionAgent) generateTitle(ctx context.Context, session *session.Sessi
698 }
724 }
699}
725}
700
726
727+func (a *sessionAgent) openrouterCost(metadata fantasy.ProviderMetadata) *float64 {
728+ openrouterMetadata, ok := metadata[openrouter.Name]
729+ if !ok {
730+ return nil
731+ }
732+
733+ opts, ok := openrouterMetadata.(*openrouter.ProviderMetadata)
734+ if !ok {
735+ return nil
736+ }
737+ return &opts.Usage.Cost
738+}
739+
701−func (a *sessionAgent) updateSessionUsage(model Model, session *session.Session, usage fantasy.Usage) {
740+func (a *sessionAgent) updateSessionUsage(model Model, session *session.Session, usage fantasy.Usage, overrideCost *float64) {
702 modelConfig := model.CatwalkCfg
741 modelConfig := model.CatwalkCfg
703 cost := modelConfig.CostPer1MInCached/1e6*float64(usage.CacheCreationTokens) +
742 cost := modelConfig.CostPer1MInCached/1e6*float64(usage.CacheCreationTokens) +
704 modelConfig.CostPer1MOutCached/1e6*float64(usage.CacheReadTokens) +
743 modelConfig.CostPer1MOutCached/1e6*float64(usage.CacheReadTokens) +
@@ -707,7 +746,12 @@ func (a *sessionAgent) updateSessionUsage(model Model, session *session.Session,
707
746
708 a.eventTokensUsed(session.ID, model, usage, cost)
747 a.eventTokensUsed(session.ID, model, usage, cost)
709
748
749+ if overrideCost != nil {
750+ session.Cost += *overrideCost
751+ } else {
710− session.Cost += cost
752+ session.Cost += cost
753+ }
754+
711 session.CompletionTokens = usage.OutputTokens + usage.CacheReadTokens
755 session.CompletionTokens = usage.OutputTokens + usage.CacheReadTokens
712 session.PromptTokens = usage.InputTokens + usage.CacheCreationTokens
756 session.PromptTokens = usage.InputTokens + usage.CacheCreationTokens
713}
757}