Skip to content

Commit 862da76

Browse files
wjiayiswjiayis
committed
feat: separate success cost from failed cost
1 parent 3dd9bd9 commit 862da76

3 files changed

Lines changed: 47 additions & 66 deletions

File tree

‎internal/models/usage.go‎

Lines changed: 20 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -9,12 +9,13 @@ import (
99
// HourlyUsage tracks cost per user, per project, per hour.
1010
// Each document represents one hour bucket of usage.
1111
type HourlyUsage struct {
12-
ID bson.ObjectID `bson:"_id"`
13-
UserID bson.ObjectID `bson:"user_id"`
14-
ProjectID string `bson:"project_id"`
15-
HourBucket bson.DateTime `bson:"hour_bucket"` // Timestamp truncated to the hour
16-
Cost float64 `bson:"cost"` // Cost in USD
17-
UpdatedAt bson.DateTime `bson:"updated_at"`
12+
ID bson.ObjectID `bson:"_id"`
13+
UserID bson.ObjectID `bson:"user_id"`
14+
ProjectID string `bson:"project_id"`
15+
HourBucket bson.DateTime `bson:"hour_bucket"` // Timestamp truncated to the hour
16+
SuccessCost float64 `bson:"success_cost"` // Cost in USD for successful requests
17+
FailedCost float64 `bson:"failed_cost"` // Cost in USD for failed requests
18+
UpdatedAt bson.DateTime `bson:"updated_at"`
1819
}
1920

2021
func (u HourlyUsage) CollectionName() string {
@@ -24,12 +25,13 @@ func (u HourlyUsage) CollectionName() string {
2425
// WeeklyUsage tracks cost per user, per project, per week.
2526
// Each document represents one week bucket of usage.
2627
type WeeklyUsage struct {
27-
ID bson.ObjectID `bson:"_id"`
28-
UserID bson.ObjectID `bson:"user_id"`
29-
ProjectID string `bson:"project_id"`
30-
WeekBucket bson.DateTime `bson:"week_bucket"` // Timestamp truncated to the week (Monday)
31-
Cost float64 `bson:"cost"` // Cost in USD
32-
UpdatedAt bson.DateTime `bson:"updated_at"`
28+
ID bson.ObjectID `bson:"_id"`
29+
UserID bson.ObjectID `bson:"user_id"`
30+
ProjectID string `bson:"project_id"`
31+
WeekBucket bson.DateTime `bson:"week_bucket"` // Timestamp truncated to the week (Monday)
32+
SuccessCost float64 `bson:"success_cost"` // Cost in USD for successful requests
33+
FailedCost float64 `bson:"failed_cost"` // Cost in USD for failed requests
34+
UpdatedAt bson.DateTime `bson:"updated_at"`
3335
}
3436

3537
func (u WeeklyUsage) CollectionName() string {
@@ -39,11 +41,12 @@ func (u WeeklyUsage) CollectionName() string {
3941
// LifetimeUsage tracks total cost per user, per project, across all time.
4042
// Each document represents the cumulative usage for a user-project pair.
4143
type LifetimeUsage struct {
42-
ID bson.ObjectID `bson:"_id"`
43-
UserID bson.ObjectID `bson:"user_id"`
44-
ProjectID string `bson:"project_id"`
45-
Cost float64 `bson:"cost"` // Total cost in USD
46-
UpdatedAt bson.DateTime `bson:"updated_at"`
44+
ID bson.ObjectID `bson:"_id"`
45+
UserID bson.ObjectID `bson:"user_id"`
46+
ProjectID string `bson:"project_id"`
47+
SuccessCost float64 `bson:"success_cost"` // Total cost in USD for successful requests
48+
FailedCost float64 `bson:"failed_cost"` // Total cost in USD for failed requests
49+
UpdatedAt bson.DateTime `bson:"updated_at"`
4750
}
4851

4952
func (u LifetimeUsage) CollectionName() string {

‎internal/services/toolkit/client/completion_v2.go‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,7 @@ func (a *AIClientV2) ChatCompletionStreamV2(ctx context.Context, callbackStream
6868
openaiChatHistory := messages
6969
inappChatHistory := AppChatHistory{}
7070
usage := UsageCost{}
71+
success := false // Track whether the request completed successfully
7172

7273
streamHandler := handler.NewStreamHandlerV2(callbackStream, conversationId, modelSlug)
7374

@@ -83,7 +84,7 @@ func (a *AIClientV2) ChatCompletionStreamV2(ctx context.Context, callbackStream
8384
// Use a detached context since the request context may be canceled
8485
trackCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
8586
defer cancel()
86-
if err := a.usageService.TrackUsage(trackCtx, userID, projectID, usage.Cost); err != nil {
87+
if err := a.usageService.TrackUsage(trackCtx, userID, projectID, usage.Cost, success); err != nil {
8788
a.logger.Error("Error while tracking usage", "error", err)
8889
}
8990
}
@@ -242,5 +243,6 @@ func (a *AIClientV2) ChatCompletionStreamV2(ctx context.Context, callbackStream
242243
}
243244
}
244245

246+
success = true
245247
return openaiChatHistory, inappChatHistory, usage, nil
246248
}

‎internal/services/usage.go‎

Lines changed: 24 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -113,43 +113,42 @@ func NewUsageService(db *db.DB, cfg *cfg.Cfg, logger *logger.Logger) *UsageServi
113113

114114
// TrackUsage increments cost for a user/project in hourly, weekly, and lifetime buckets.
115115
// Uses upsert to create or update the usage records atomically.
116-
func (s *UsageService) TrackUsage(ctx context.Context, userID bson.ObjectID, projectID string, cost float64) error {
116+
// The success parameter indicates whether the request completed successfully.
117+
// We will be charging only for successful requests, but we track failed requests for monitoring.
118+
func (s *UsageService) TrackUsage(ctx context.Context, userID bson.ObjectID, projectID string, cost float64, success bool) error {
117119
if cost == 0 {
118120
return nil
119121
}
120122

121123
now := time.Now()
122124

123125
// Track hourly usage
124-
if err := s.trackHourlyUsage(ctx, userID, projectID, cost, now); err != nil {
126+
if err := s.trackHourlyUsage(ctx, userID, projectID, cost, success, now); err != nil {
125127
return err
126128
}
127129

128130
// Track weekly usage
129-
if err := s.trackWeeklyUsage(ctx, userID, projectID, cost, now); err != nil {
131+
if err := s.trackWeeklyUsage(ctx, userID, projectID, cost, success, now); err != nil {
130132
return err
131133
}
132134

133135
// Track lifetime usage
134-
if err := s.trackLifetimeUsage(ctx, userID, projectID, cost, now); err != nil {
136+
if err := s.trackLifetimeUsage(ctx, userID, projectID, cost, success, now); err != nil {
135137
return err
136138
}
137139

138140
return nil
139141
}
140142

141-
func (s *UsageService) trackHourlyUsage(ctx context.Context, userID bson.ObjectID, projectID string, cost float64, now time.Time) error {
142-
hourBucket := models.TruncateToHour(now)
143-
144-
filter := bson.M{
145-
"user_id": userID,
146-
"project_id": projectID,
147-
"hour_bucket": bson.NewDateTimeFromTime(hourBucket),
143+
func (s *UsageService) upsertUsage(ctx context.Context, collection *mongo.Collection, filter bson.M, cost float64, success bool, now time.Time) error {
144+
costField := "failed_cost"
145+
if success {
146+
costField = "success_cost"
148147
}
149148

150149
update := bson.M{
151150
"$inc": bson.M{
152-
"cost": cost,
151+
costField: cost,
153152
},
154153
"$set": bson.M{
155154
"updated_at": bson.NewDateTimeFromTime(now),
@@ -160,55 +159,32 @@ func (s *UsageService) trackHourlyUsage(ctx context.Context, userID bson.ObjectI
160159
}
161160

162161
opts := options.UpdateOne().SetUpsert(true)
163-
_, err := s.hourlyCollection.UpdateOne(ctx, filter, update, opts)
162+
_, err := collection.UpdateOne(ctx, filter, update, opts)
164163
return err
165164
}
166165

167-
func (s *UsageService) trackWeeklyUsage(ctx context.Context, userID bson.ObjectID, projectID string, cost float64, now time.Time) error {
168-
weekBucket := models.TruncateToWeek(now)
169-
166+
func (s *UsageService) trackHourlyUsage(ctx context.Context, userID bson.ObjectID, projectID string, cost float64, success bool, now time.Time) error {
170167
filter := bson.M{
171168
"user_id": userID,
172169
"project_id": projectID,
173-
"week_bucket": bson.NewDateTimeFromTime(weekBucket),
170+
"hour_bucket": bson.NewDateTimeFromTime(models.TruncateToHour(now)),
174171
}
172+
return s.upsertUsage(ctx, s.hourlyCollection, filter, cost, success, now)
173+
}
175174

176-
update := bson.M{
177-
"$inc": bson.M{
178-
"cost": cost,
179-
},
180-
"$set": bson.M{
181-
"updated_at": bson.NewDateTimeFromTime(now),
182-
},
183-
"$setOnInsert": bson.M{
184-
"_id": bson.NewObjectID(),
185-
},
175+
func (s *UsageService) trackWeeklyUsage(ctx context.Context, userID bson.ObjectID, projectID string, cost float64, success bool, now time.Time) error {
176+
filter := bson.M{
177+
"user_id": userID,
178+
"project_id": projectID,
179+
"week_bucket": bson.NewDateTimeFromTime(models.TruncateToWeek(now)),
186180
}
187-
188-
opts := options.UpdateOne().SetUpsert(true)
189-
_, err := s.weeklyCollection.UpdateOne(ctx, filter, update, opts)
190-
return err
181+
return s.upsertUsage(ctx, s.weeklyCollection, filter, cost, success, now)
191182
}
192183

193-
func (s *UsageService) trackLifetimeUsage(ctx context.Context, userID bson.ObjectID, projectID string, cost float64, now time.Time) error {
184+
func (s *UsageService) trackLifetimeUsage(ctx context.Context, userID bson.ObjectID, projectID string, cost float64, success bool, now time.Time) error {
194185
filter := bson.M{
195186
"user_id": userID,
196187
"project_id": projectID,
197188
}
198-
199-
update := bson.M{
200-
"$inc": bson.M{
201-
"cost": cost,
202-
},
203-
"$set": bson.M{
204-
"updated_at": bson.NewDateTimeFromTime(now),
205-
},
206-
"$setOnInsert": bson.M{
207-
"_id": bson.NewObjectID(),
208-
},
209-
}
210-
211-
opts := options.UpdateOne().SetUpsert(true)
212-
_, err := s.lifetimeCollection.UpdateOne(ctx, filter, update, opts)
213-
return err
189+
return s.upsertUsage(ctx, s.lifetimeCollection, filter, cost, success, now)
214190
}

0 commit comments

Comments
 (0)