Function TestChatCostSummary_UnpricedMessages

coderd/exp_chats_test.go:10427–10481

(t *testing.T)

}

func TestChatCostSummary_UnpricedMessages(t *testing.T) {
	t.Parallel()

	ctx := testutil.Context(t, testutil.WaitLong)
	client, db := newChatClientWithDatabase(t)
	firstUser := coderdtest.CreateFirstUser(t, client.Client)
	modelConfig := createChatModelConfig(t, client)

	chat := dbgen.Chat(t, db, database.Chat{
		OrganizationID:    firstUser.OrganizationID,
		OwnerID:           firstUser.UserID,
		LastModelConfigID: modelConfig.ID,
		Title:             "unpriced test",
	})

	pricedMessage := dbgen.ChatMessage(t, db, database.ChatMessage{
		ChatID:          chat.ID,
		ModelConfigID:   uuid.NullUUID{UUID: modelConfig.ID, Valid: true},
		Role:            database.ChatMessageRoleAssistant,
		InputTokens:     sql.NullInt64{Int64: 100, Valid: true},
		OutputTokens:    sql.NullInt64{Int64: 50, Valid: true},
		TotalCostMicros: sql.NullInt64{Int64: 500, Valid: true},
	})

	unpricedMessage := dbgen.ChatMessage(t, db, database.ChatMessage{
		ChatID:        chat.ID,
		ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true},
		Role:          database.ChatMessageRoleAssistant,
		InputTokens:   sql.NullInt64{Int64: 200, Valid: true},
		OutputTokens:  sql.NullInt64{Int64: 75, Valid: true},
	})

	earliestCreatedAt := pricedMessage.CreatedAt
	latestCreatedAt := pricedMessage.CreatedAt
	if unpricedMessage.CreatedAt.Before(earliestCreatedAt) {
		earliestCreatedAt = unpricedMessage.CreatedAt
	}
	if unpricedMessage.CreatedAt.After(latestCreatedAt) {
		latestCreatedAt = unpricedMessage.CreatedAt
	}
	options := codersdk.ChatCostSummaryOptions{
		// Pad the DB-assigned timestamps to avoid time boundary flakes.
		StartDate: earliestCreatedAt.Add(-time.Minute),
		EndDate:   latestCreatedAt.Add(time.Minute),
	}

	summary, err := client.GetChatCostSummary(ctx, "me", options)
	require.NoError(t, err)

	require.Equal(t, int64(500), summary.TotalCostMicros)
	require.Equal(t, int64(1), summary.PricedMessageCount)
	require.Equal(t, int64(1), summary.UnpricedMessageCount)
	require.Equal(t, int64(300), summary.TotalInputTokens)
	require.Equal(t, int64(125), summary.TotalOutputTokens)
}

func requireChatModelPricing(
	t *testing.T,