diff --git a/internal/chat/service_usage.go b/internal/chat/service_usage.go new file mode 100644 index 0000000..884c018 --- /dev/null +++ b/internal/chat/service_usage.go @@ -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) +} diff --git a/internal/chat/service_usage_test.go b/internal/chat/service_usage_test.go new file mode 100644 index 0000000..af09f17 --- /dev/null +++ b/internal/chat/service_usage_test.go @@ -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) +}