feat(heatmap): add HeatmapViewModel with quantile-based color levels and three view modes
- Daily/Weekly/Cumulative data aggregation from OCGO records - Quantile (P25/P50/P75/P95) color level mapping with nearest-rank - Token = inputTokens + cacheReadTokens + reasoningTokens + outputTokens - 365-day rolling window, 52-week weekly view, cumulative running sum - Hilt injected, Dispatchers.Default for heavy ops, try-catch for safety
This commit is contained in:
parent
7f2ca18d9d
commit
b8479eebf3
322
app/src/main/java/com/rainy/token/ui/heatmap/HeatmapViewModel.kt
Normal file
322
app/src/main/java/com/rainy/token/ui/heatmap/HeatmapViewModel.kt
Normal file
@ -0,0 +1,322 @@
|
|||||||
|
package com.rainy.token.ui.heatmap
|
||||||
|
|
||||||
|
import androidx.lifecycle.ViewModel
|
||||||
|
import androidx.lifecycle.viewModelScope
|
||||||
|
import com.rainy.token.data.local.UsageCache
|
||||||
|
import com.rainy.token.data.local.UsageRecord
|
||||||
|
import com.rainy.token.data.repository.CredentialRepository
|
||||||
|
import com.rainy.token.domain.model.Credential
|
||||||
|
import com.rainy.token.domain.service.ServiceType
|
||||||
|
import dagger.hilt.android.lifecycle.HiltViewModel
|
||||||
|
import kotlinx.coroutines.Dispatchers
|
||||||
|
import kotlinx.coroutines.flow.MutableStateFlow
|
||||||
|
import kotlinx.coroutines.flow.StateFlow
|
||||||
|
import kotlinx.coroutines.flow.asStateFlow
|
||||||
|
import kotlinx.coroutines.flow.update
|
||||||
|
import kotlinx.coroutines.launch
|
||||||
|
import kotlinx.coroutines.withContext
|
||||||
|
import javax.inject.Inject
|
||||||
|
import javax.inject.Provider
|
||||||
|
|
||||||
|
// ── 公共数据类 ──────────────────────────────────────────────
|
||||||
|
|
||||||
|
/** 热力图视图模式 */
|
||||||
|
enum class HeatmapViewMode { DAILY, WEEKLY, CUMULATIVE }
|
||||||
|
|
||||||
|
/** 每日 / 累计视图的单格数据 */
|
||||||
|
data class HeatmapDayData(
|
||||||
|
val dayTs: Long, // UTC 日期时间戳(即天分桶的时间戳)
|
||||||
|
val tokens: Long, // 当天 token 数(DAILY)或累计 token 数(CUMULATIVE)
|
||||||
|
val level: Int, // 0-5 颜色等级
|
||||||
|
)
|
||||||
|
|
||||||
|
/** 每周视图的单格数据 */
|
||||||
|
data class HeatmapWeekData(
|
||||||
|
val weekStartTs: Long, // 该周起始日时间戳
|
||||||
|
val tokens: Long, // 该周 7 天累计 token 数
|
||||||
|
val level: Int, // 0-5 颜色等级
|
||||||
|
)
|
||||||
|
|
||||||
|
/** UI 状态 */
|
||||||
|
data class HeatmapUiState(
|
||||||
|
val loading: Boolean = true,
|
||||||
|
val viewMode: HeatmapViewMode = HeatmapViewMode.DAILY,
|
||||||
|
val dailyData: List<HeatmapDayData> = emptyList(),
|
||||||
|
val weeklyData: List<HeatmapWeekData> = emptyList(),
|
||||||
|
val cumulativeData: List<HeatmapDayData> = emptyList(),
|
||||||
|
val colorLevels: IntArray = IntArray(6), // [0, p25, p50, p75, p95, max]
|
||||||
|
) {
|
||||||
|
override fun equals(other: Any?): Boolean {
|
||||||
|
if (this === other) return true
|
||||||
|
if (other !is HeatmapUiState) return false
|
||||||
|
return loading == other.loading &&
|
||||||
|
viewMode == other.viewMode &&
|
||||||
|
dailyData == other.dailyData &&
|
||||||
|
weeklyData == other.weeklyData &&
|
||||||
|
cumulativeData == other.cumulativeData &&
|
||||||
|
colorLevels.contentEquals(other.colorLevels)
|
||||||
|
}
|
||||||
|
|
||||||
|
override fun hashCode(): Int {
|
||||||
|
var result = loading.hashCode()
|
||||||
|
result = 31 * result + viewMode.hashCode()
|
||||||
|
result = 31 * result + dailyData.hashCode()
|
||||||
|
result = 31 * result + weeklyData.hashCode()
|
||||||
|
result = 31 * result + cumulativeData.hashCode()
|
||||||
|
result = 31 * result + colorLevels.contentHashCode()
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 分位阈值与颜色等级映射。
|
||||||
|
*
|
||||||
|
* thresholds[0..5] 存 Long 阈值:
|
||||||
|
* - Level 0: value == 0
|
||||||
|
* - Level 1: 0 < value ≤ P25
|
||||||
|
* - Level 2: P25 < value ≤ P50
|
||||||
|
* - Level 3: P50 < value ≤ P75
|
||||||
|
* - Level 4: P75 < value ≤ P95
|
||||||
|
* - Level 5: value > P95
|
||||||
|
*
|
||||||
|
* 边界情况(非零天数 < 2 或所有非零值相同)时,所有非零值设为 Level 3。
|
||||||
|
*/
|
||||||
|
internal data class QuantileLevels(val thresholds: LongArray) {
|
||||||
|
fun getColorLevel(value: Long): Int {
|
||||||
|
if (value <= 0L) return 0
|
||||||
|
for (level in 1..5) {
|
||||||
|
if (value <= thresholds[level]) return level
|
||||||
|
}
|
||||||
|
return 5
|
||||||
|
}
|
||||||
|
|
||||||
|
fun toIntArray(): IntArray = IntArray(6) { idx ->
|
||||||
|
val v = thresholds[idx]
|
||||||
|
if (v > Int.MAX_VALUE) Int.MAX_VALUE else v.toInt()
|
||||||
|
}
|
||||||
|
|
||||||
|
override fun equals(other: Any?): Boolean {
|
||||||
|
if (this === other) return true
|
||||||
|
if (other !is QuantileLevels) return false
|
||||||
|
return thresholds.contentEquals(other.thresholds)
|
||||||
|
}
|
||||||
|
|
||||||
|
override fun hashCode(): Int = thresholds.contentHashCode()
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── ViewModel ──────────────────────────────────────────────
|
||||||
|
|
||||||
|
@HiltViewModel
|
||||||
|
class HeatmapViewModel @Inject constructor(
|
||||||
|
private val cacheProvider: Provider<UsageCache>,
|
||||||
|
private val credentialRepository: CredentialRepository,
|
||||||
|
) : ViewModel() {
|
||||||
|
|
||||||
|
private val _uiState = MutableStateFlow(HeatmapUiState())
|
||||||
|
val uiState: StateFlow<HeatmapUiState> = _uiState.asStateFlow()
|
||||||
|
|
||||||
|
private var loaded = false
|
||||||
|
private var allViewsResult: AllViewsResult? = null
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Composable 层在 LaunchedEffect 中调用一次即可。
|
||||||
|
* 三种视图数据全部在此次加载中计算好,切换视图不需要重新加载。
|
||||||
|
*/
|
||||||
|
fun load() {
|
||||||
|
if (loaded) return
|
||||||
|
loaded = true
|
||||||
|
viewModelScope.launch { loadInternal() }
|
||||||
|
}
|
||||||
|
|
||||||
|
/** 切换视图模式(仅改 State,不重新加载数据) */
|
||||||
|
fun setViewMode(mode: HeatmapViewMode) {
|
||||||
|
val result = allViewsResult
|
||||||
|
val levels = result?.let {
|
||||||
|
when (mode) {
|
||||||
|
HeatmapViewMode.DAILY -> it.dailyLevels
|
||||||
|
HeatmapViewMode.WEEKLY -> it.weeklyLevels
|
||||||
|
HeatmapViewMode.CUMULATIVE -> it.cumulativeLevels
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_uiState.update { it.copy(viewMode = mode, colorLevels = levels?.toIntArray() ?: it.colorLevels) }
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── 内部逻辑 ────────────────────────────────────────────
|
||||||
|
|
||||||
|
private suspend fun workspaceId(): String? {
|
||||||
|
val c = credentialRepository.get(ServiceType.OPENCODE_GO)
|
||||||
|
return (c as? Credential.SessionCredential)?.workspaceId?.takeIf { it.isNotBlank() }
|
||||||
|
}
|
||||||
|
|
||||||
|
private suspend fun loadInternal() {
|
||||||
|
try {
|
||||||
|
val wid = workspaceId() ?: run {
|
||||||
|
_uiState.update { it.copy(loading = false) }
|
||||||
|
return
|
||||||
|
}
|
||||||
|
val cache = cacheProvider.get()
|
||||||
|
|
||||||
|
// 所有重操作放在 Dispatchers.Default 上
|
||||||
|
val result = withContext(Dispatchers.Default) {
|
||||||
|
val records = cache.getRecords(wid)
|
||||||
|
computeAllViews(records)
|
||||||
|
}
|
||||||
|
|
||||||
|
allViewsResult = result
|
||||||
|
|
||||||
|
_uiState.update {
|
||||||
|
it.copy(
|
||||||
|
loading = false,
|
||||||
|
dailyData = result.dailyData,
|
||||||
|
weeklyData = result.weeklyData,
|
||||||
|
cumulativeData = result.cumulativeData,
|
||||||
|
colorLevels = result.dailyLevels.toIntArray(),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
} catch (e: Exception) {
|
||||||
|
_uiState.update { it.copy(loading = false) }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── 聚合计算 ────────────────────────────────────────────
|
||||||
|
|
||||||
|
companion object {
|
||||||
|
private const val DAY_MS = 86_400_000L
|
||||||
|
private const val RANGE_DAYS = 365
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 从全部记录计算三种视图数据 + 颜色等级阈值。
|
||||||
|
* 返回 (dailyData, weeklyData, cumulativeData, colorLevels)
|
||||||
|
*/
|
||||||
|
internal fun computeAllViews(records: List<UsageRecord>): AllViewsResult {
|
||||||
|
val now = System.currentTimeMillis()
|
||||||
|
|
||||||
|
// 1. 按天分桶聚合 token 数
|
||||||
|
val todayTs = now / DAY_MS * DAY_MS
|
||||||
|
val rangeStartTs = todayTs - (RANGE_DAYS - 1) * DAY_MS
|
||||||
|
|
||||||
|
// 聚合每一天的 token
|
||||||
|
val dayTokenMap = HashMap<Long, Long>()
|
||||||
|
for (r in records) {
|
||||||
|
val dayTs = r.timeCreated / DAY_MS * DAY_MS
|
||||||
|
if (dayTs < rangeStartTs || dayTs > todayTs) continue
|
||||||
|
val tokens = r.inputTokens + r.cacheReadTokens + r.reasoningTokens + r.outputTokens
|
||||||
|
dayTokenMap.merge(dayTs, tokens, Long::plus)
|
||||||
|
}
|
||||||
|
|
||||||
|
// 2. 构建完整的每日数据列表(365 天,从 rangeStartTs 到 todayTs)
|
||||||
|
val dailyRaw = ArrayList<HeatmapDayData>(RANGE_DAYS)
|
||||||
|
for (i in 0 until RANGE_DAYS) {
|
||||||
|
val ts = rangeStartTs + i * DAY_MS
|
||||||
|
val tokens = dayTokenMap[ts] ?: 0L
|
||||||
|
dailyRaw.add(HeatmapDayData(dayTs = ts, tokens = tokens, level = 0))
|
||||||
|
}
|
||||||
|
|
||||||
|
// 3. 计算每日视图颜色等级
|
||||||
|
val dailyLevels = computeQuantileLevels(dailyRaw.map { it.tokens })
|
||||||
|
val dailyData = dailyRaw.map { it.copy(level = dailyLevels.getColorLevel(it.tokens)) }
|
||||||
|
|
||||||
|
// 4. 每周视图:从最早日数据开始,每 7 天一组
|
||||||
|
val (weeklyData, weeklyLevels) = buildWeeklyData(dailyRaw)
|
||||||
|
|
||||||
|
// 5. 累计视图:从第一天到当天的累计总和
|
||||||
|
var cumulativeSum = 0L
|
||||||
|
val cumulativeRaw = ArrayList<HeatmapDayData>(RANGE_DAYS)
|
||||||
|
for (d in dailyRaw) {
|
||||||
|
cumulativeSum += d.tokens
|
||||||
|
cumulativeRaw.add(HeatmapDayData(dayTs = d.dayTs, tokens = cumulativeSum, level = 0))
|
||||||
|
}
|
||||||
|
val cumulativeLevels = computeQuantileLevels(cumulativeRaw.map { it.tokens })
|
||||||
|
val cumulativeData = cumulativeRaw.map { it.copy(level = cumulativeLevels.getColorLevel(it.tokens)) }
|
||||||
|
|
||||||
|
// 保存三种视图各自的分位阈值,供切换视图时更新图例
|
||||||
|
return AllViewsResult(dailyData, weeklyData, cumulativeData, dailyLevels, weeklyLevels, cumulativeLevels)
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 构建每周视图数据。
|
||||||
|
* 从 rangeStartTs 对应的那天开始,每 7 天为一组。
|
||||||
|
*/
|
||||||
|
private fun buildWeeklyData(dailyRaw: List<HeatmapDayData>): Pair<List<HeatmapWeekData>, QuantileLevels> {
|
||||||
|
val weeks = ArrayList<HeatmapWeekData>()
|
||||||
|
var i = 0
|
||||||
|
while (i < 52 * 7 && i < dailyRaw.size) {
|
||||||
|
val weekEnd = minOf(i + 7, dailyRaw.size)
|
||||||
|
var weekTokens = 0L
|
||||||
|
for (j in i until weekEnd) {
|
||||||
|
weekTokens += dailyRaw[j].tokens
|
||||||
|
}
|
||||||
|
weeks.add(HeatmapWeekData(
|
||||||
|
weekStartTs = dailyRaw[i].dayTs,
|
||||||
|
tokens = weekTokens,
|
||||||
|
level = 0,
|
||||||
|
))
|
||||||
|
i += 7
|
||||||
|
}
|
||||||
|
|
||||||
|
// 计算每周视图颜色等级
|
||||||
|
val weekLevels = computeQuantileLevels(weeks.map { it.tokens })
|
||||||
|
return weeks.map { it.copy(level = weekLevels.getColorLevel(it.tokens)) } to weekLevels
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── 分位数颜色等级计算 ──────────────────────────────
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 使用 nearest-rank 分位数法计算 P25/P50/P75/P95 阈值。
|
||||||
|
*
|
||||||
|
* 边界情况:
|
||||||
|
* - 非零天数 < 2 → 全部设为 Level 3
|
||||||
|
* - 所有非零值相同 → 全部设为 Level 3
|
||||||
|
*
|
||||||
|
* 返回 [QuantileLevels],内部 thresholds 为 LongArray(6):
|
||||||
|
* index 0 = 0(Level 0),index 1-5 分别存 P25/P50/P75/P95/MaxValue。
|
||||||
|
*/
|
||||||
|
internal fun computeQuantileLevels(values: List<Long>): QuantileLevels {
|
||||||
|
// 收集所有非零值
|
||||||
|
val nonZero = values.filter { it > 0L }
|
||||||
|
|
||||||
|
// 边界:非零值 < 2 或全部相同 → 全部设为 Level 3
|
||||||
|
if (nonZero.size < 2 || nonZero.toSet().size == 1) {
|
||||||
|
// thresholds[3] = Long.MAX_VALUE 使得 0 < value ≤ MAX → Level 3
|
||||||
|
val t = LongArray(6)
|
||||||
|
t[3] = Long.MAX_VALUE
|
||||||
|
return QuantileLevels(t)
|
||||||
|
}
|
||||||
|
|
||||||
|
val sorted = nonZero.sorted()
|
||||||
|
val n = sorted.size
|
||||||
|
|
||||||
|
val p25 = nearestRank(sorted, n, 0.25)
|
||||||
|
val p50 = nearestRank(sorted, n, 0.50)
|
||||||
|
val p75 = nearestRank(sorted, n, 0.75)
|
||||||
|
val p95 = nearestRank(sorted, n, 0.95)
|
||||||
|
|
||||||
|
// LongArray(6): [0, p25, p50, p75, p95, MaxValue]
|
||||||
|
val thresholds = LongArray(6)
|
||||||
|
thresholds[0] = 0L
|
||||||
|
thresholds[1] = p25
|
||||||
|
thresholds[2] = p50
|
||||||
|
thresholds[3] = p75
|
||||||
|
thresholds[4] = p95
|
||||||
|
thresholds[5] = Long.MAX_VALUE
|
||||||
|
return QuantileLevels(thresholds)
|
||||||
|
}
|
||||||
|
|
||||||
|
private fun nearestRank(sorted: List<Long>, n: Int, percentile: Double): Long {
|
||||||
|
val rank = (Math.ceil(n * percentile).toInt() - 1).coerceIn(0, n - 1)
|
||||||
|
return sorted[rank]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── 辅助类型 ──────────────────────────────────────────────
|
||||||
|
|
||||||
|
/** computeAllViews 返回值 */
|
||||||
|
internal data class AllViewsResult(
|
||||||
|
val dailyData: List<HeatmapDayData>,
|
||||||
|
val weeklyData: List<HeatmapWeekData>,
|
||||||
|
val cumulativeData: List<HeatmapDayData>,
|
||||||
|
val dailyLevels: QuantileLevels,
|
||||||
|
val weeklyLevels: QuantileLevels,
|
||||||
|
val cumulativeLevels: QuantileLevels,
|
||||||
|
)
|
||||||
Loading…
Reference in New Issue
Block a user