feat(chat): UsageService for cost/token summarization

This commit is contained in:
2026-05-04 17:42:36 +08:00
parent c42887498b
commit 8b5dc70146
2 changed files with 133 additions and 0 deletions
+35
View File
@@ -0,0 +1,35 @@
package chat
import (
"context"
"time"
"github.com/google/uuid"
)
type usageService struct{ repo Repository }
var _ UsageService = (*usageService)(nil)
// NewUsageService 构造 llm_usage 写入与汇总查询服务。
func NewUsageService(repo Repository) UsageService { return &usageService{repo: repo} }
func (s *usageService) Insert(ctx context.Context, in UsageRecord) error {
return s.repo.InsertUsage(ctx, in)
}
func (s *usageService) SummarizeByUser(ctx context.Context, from, to time.Time) ([]UserUsageRow, error) {
return s.repo.SummarizeUsageByUser(ctx, from, to)
}
func (s *usageService) SummarizeByModel(ctx context.Context, from, to time.Time) ([]ModelUsageRow, error) {
return s.repo.SummarizeUsageByModel(ctx, from, to)
}
func (s *usageService) SummarizeByDay(ctx context.Context, from, to time.Time) ([]DayUsageRow, error) {
return s.repo.SummarizeUsageByDay(ctx, from, to)
}
func (s *usageService) UserDaily(ctx context.Context, userID uuid.UUID, from, to time.Time) ([]DayUsageRow, error) {
return s.repo.SummarizeUserDaily(ctx, userID, from, to)
}
+98
View File
@@ -0,0 +1,98 @@
package chat
import (
"context"
"testing"
"time"
"github.com/google/uuid"
"github.com/stretchr/testify/require"
)
// usageRepo is a tiny stub that records what was called and returns canned values.
// It satisfies only the UsageService-relevant subset of Repository; unused methods
// inherit panic stubs from fakeRepo so we don't need to declare them here.
type usageRepo struct {
*fakeRepo
insertedUsage []UsageRecord
byUserCalls []timeRange
byModelCalls []timeRange
byDayCalls []timeRange
dailyCalls []userTimeRange
usersOut []UserUsageRow
modelsOut []ModelUsageRow
daysOut []DayUsageRow
}
type timeRange struct{ from, to time.Time }
type userTimeRange struct {
user uuid.UUID
from time.Time
to time.Time
}
func (r *usageRepo) InsertUsage(_ context.Context, rec UsageRecord) error {
r.insertedUsage = append(r.insertedUsage, rec)
return nil
}
func (r *usageRepo) SummarizeUsageByUser(_ context.Context, from, to time.Time) ([]UserUsageRow, error) {
r.byUserCalls = append(r.byUserCalls, timeRange{from, to})
return r.usersOut, nil
}
func (r *usageRepo) SummarizeUsageByModel(_ context.Context, from, to time.Time) ([]ModelUsageRow, error) {
r.byModelCalls = append(r.byModelCalls, timeRange{from, to})
return r.modelsOut, nil
}
func (r *usageRepo) SummarizeUsageByDay(_ context.Context, from, to time.Time) ([]DayUsageRow, error) {
r.byDayCalls = append(r.byDayCalls, timeRange{from, to})
return r.daysOut, nil
}
func (r *usageRepo) SummarizeUserDaily(_ context.Context, uid uuid.UUID, from, to time.Time) ([]DayUsageRow, error) {
r.dailyCalls = append(r.dailyCalls, userTimeRange{uid, from, to})
return r.daysOut, nil
}
func TestUsageService_DelegatesToRepo(t *testing.T) {
repo := &usageRepo{fakeRepo: newFakeRepo()}
svc := NewUsageService(repo)
from := time.Date(2026, 5, 1, 0, 0, 0, 0, time.UTC)
to := from.Add(24 * time.Hour)
uid := uuid.Must(uuid.NewV7())
rec := UsageRecord{
ConversationID: uuid.Must(uuid.NewV7()),
MessageID: 42,
UserID: uid,
PromptTokens: 100, CompletionTokens: 50, ThinkingTokens: 0,
CostUSD: 0.001,
}
require.NoError(t, svc.Insert(context.Background(), rec))
require.Len(t, repo.insertedUsage, 1)
require.Equal(t, rec, repo.insertedUsage[0])
repo.usersOut = []UserUsageRow{{UserID: uid, PromptTokens: 100}}
users, err := svc.SummarizeByUser(context.Background(), from, to)
require.NoError(t, err)
require.Equal(t, repo.usersOut, users)
require.Equal(t, []timeRange{{from, to}}, repo.byUserCalls)
repo.modelsOut = []ModelUsageRow{{ModelID: uuid.Must(uuid.NewV7()), CostUSD: 1.5}}
models, err := svc.SummarizeByModel(context.Background(), from, to)
require.NoError(t, err)
require.Equal(t, repo.modelsOut, models)
require.Equal(t, []timeRange{{from, to}}, repo.byModelCalls)
repo.daysOut = []DayUsageRow{{Day: from, PromptTokens: 200}}
days, err := svc.SummarizeByDay(context.Background(), from, to)
require.NoError(t, err)
require.Equal(t, repo.daysOut, days)
require.Equal(t, []timeRange{{from, to}}, repo.byDayCalls)
dailyOut, err := svc.UserDaily(context.Background(), uid, from, to)
require.NoError(t, err)
require.Equal(t, repo.daysOut, dailyOut)
require.Equal(t, []userTimeRange{{uid, from, to}}, repo.dailyCalls)
}