diff --git a/app/src/main/java/dev/kaizen/app/chat/ChatComponents.kt b/app/src/main/java/dev/kaizen/app/chat/ChatComponents.kt index 8a68160..9101046 100644 --- a/app/src/main/java/dev/kaizen/app/chat/ChatComponents.kt +++ b/app/src/main/java/dev/kaizen/app/chat/ChatComponents.kt @@ -68,6 +68,7 @@ import dev.kaizen.app.net.Attachment import dev.kaizen.app.ui.effect.GlassSurface import dev.kaizen.app.ui.effect.GlassTiers import dev.kaizen.app.ui.effect.KaizenShadows +import dev.kaizen.app.ui.effect.kaizenShadow import dev.kaizen.app.ui.motion.Durations import dev.kaizen.app.ui.shape.KaizenShapes import dev.kaizen.app.ui.theme.LocalKaizenAccent @@ -87,6 +88,7 @@ import androidx.compose.material.icons.rounded.PictureAsPdf import androidx.compose.material.icons.rounded.VideoFile import androidx.compose.runtime.mutableStateOf import androidx.compose.runtime.setValue +import androidx.compose.material.icons.rounded.Bolt import androidx.compose.material.icons.rounded.ExpandMore import androidx.compose.material.icons.rounded.Psychology import androidx.compose.material.icons.rounded.Tune @@ -298,6 +300,41 @@ private fun SamplingRow(label: String, param: SamplingParam, min: Float, max: Fl ) } +@Composable +fun TokenBadge(used: Int, contextLength: Int, modifier: Modifier = Modifier) { + val cs = MaterialTheme.colorScheme + val isDark = isSystemInDarkTheme() + val accent = LocalKaizenAccent.current + + fun formatK(n: Int): String = when { + n >= 1_000_000 -> "${n / 1_000_000}M" + n >= 1_000 -> "${n / 1_000}k" + else -> "$n" + } + + GlassSurface( + modifier = modifier, + shape = KaizenShapes.pill, + tintAlpha = GlassTiers.pill(isDark), + shadowStack = KaizenShadows.level1, + cornerRadius = 16.dp, + ) { + Row( + Modifier.padding(horizontal = 10.dp, vertical = 7.dp), + verticalAlignment = Alignment.CenterVertically, + ) { + Icon(Icons.Rounded.Bolt, null, tint = accent.primary, modifier = Modifier.size(14.dp)) + Spacer(Modifier.width(4.dp)) + Text( + "${formatK(used)} / ${formatK(contextLength)}", + color = cs.onSurfaceVariant, + fontSize = 12.sp, + fontWeight = FontWeight.Medium, + ) + } + } +} + @Composable fun MessageRow(message: Message) { val cs = MaterialTheme.colorScheme @@ -308,22 +345,46 @@ fun MessageRow(message: Message) { if (message.role == Role.User) { Row(Modifier.fillMaxWidth(), horizontalArrangement = Arrangement.End) { val shape = KaizenShapes.lg - val borderBrush = remember(accent) { Brush.verticalGradient(listOf(accent.primary.copy(alpha = 0.45f), accent.primaryDeep.copy(alpha = 0.15f))) } - val bgBrush = remember(accent) { Brush.verticalGradient(listOf(accent.primary.copy(alpha = 0.18f), accent.primaryDeep.copy(alpha = 0.08f))) } + val bgBrush = remember(accent, isDark) { + Brush.verticalGradient( + listOf( + if (isDark) accent.primary.copy(alpha = 0.28f) else accent.primary.copy(alpha = 0.16f), + if (isDark) accent.primaryDeep.copy(alpha = 0.22f) else accent.primaryDeep.copy(alpha = 0.10f), + ) + ) + } + val borderBrush = remember(accent, isDark) { + Brush.verticalGradient( + listOf( + accent.primary.copy(alpha = if (isDark) 0.55f else 0.40f), + accent.primaryDeep.copy(alpha = if (isDark) 0.15f else 0.10f), + ) + ) + } + val innerHighlight = remember(accent, isDark) { + Brush.verticalGradient( + listOf( + accent.primary.copy(alpha = if (isDark) 0.12f else 0.08f), + Color.Transparent, + ) + ) + } Box( Modifier .widthIn(max = 300.dp) + .kaizenShadow(stack = KaizenShadows.level1, cornerRadius = 20.dp) .clip(shape) .background(bgBrush) + .background(innerHighlight) .border(1.2.dp, borderBrush, shape), ) { Column { if (hasAttachments) AttachmentGrid(message.attachments, Modifier.padding(6.dp)) if (message.content.isNotBlank()) { Text( - message.content, color = cs.onBackground, fontSize = 16.sp, lineHeight = 22.sp, - modifier = Modifier.padding(start = 16.dp, end = 16.dp, top = if (hasAttachments) 4.dp else 11.dp, bottom = 11.dp), + message.content, color = cs.onBackground, fontSize = 16.sp, lineHeight = 26.sp, + modifier = Modifier.padding(start = 16.dp, end = 16.dp, top = if (hasAttachments) 4.dp else 13.dp, bottom = 13.dp), ) } } diff --git a/app/src/main/java/dev/kaizen/app/chat/ChatScreen.kt b/app/src/main/java/dev/kaizen/app/chat/ChatScreen.kt index 5fa7d72..4306b07 100644 --- a/app/src/main/java/dev/kaizen/app/chat/ChatScreen.kt +++ b/app/src/main/java/dev/kaizen/app/chat/ChatScreen.kt @@ -24,6 +24,8 @@ import androidx.compose.foundation.lazy.LazyColumn import androidx.compose.foundation.lazy.items import androidx.compose.foundation.lazy.rememberLazyListState import androidx.compose.material.icons.Icons +import androidx.compose.material.icons.rounded.KeyboardArrowDown +import androidx.compose.material.icons.rounded.KeyboardArrowUp import androidx.compose.material.icons.rounded.Menu import androidx.compose.material3.DrawerValue import androidx.compose.material3.Icon @@ -100,6 +102,7 @@ fun ChatScreen( var activeMode by remember { mutableStateOf(null) } var reasoningPreset by remember { mutableStateOf(ReasoningPreset.Standard) } var samplingPrefs by remember { mutableStateOf(SamplingPrefs()) } + var usedTokens by remember { mutableStateOf(0) } // Navigation State var currentScreen by remember { mutableStateOf(AppScreen.Chat) } @@ -362,6 +365,9 @@ fun ChatScreen( query = state.query, sources = state.sources, ) + if (state.inputTokens > 0 || state.outputTokens > 0) { + usedTokens = state.inputTokens + state.outputTokens + } } val idx = messages.indexOfLast { it.id == assistantId } val finalMsg = if (idx >= 0) messages[idx] else null @@ -538,6 +544,53 @@ fun ChatScreen( ) { items(messages, key = { it.id }) { message -> MessageRow(message) } } + + // Scroll buttons + val showScrollUp by remember { derivedStateOf { listState.firstVisibleItemIndex > 2 } } + val showScrollDown by remember { derivedStateOf { + val info = listState.layoutInfo + info.visibleItemsInfo.lastOrNull()?.index != info.totalItemsCount - 1 + }} + + if (showScrollUp || showScrollDown) { + Column( + modifier = Modifier + .align(Alignment.CenterEnd) + .padding(end = 8.dp), + verticalArrangement = Arrangement.spacedBy(8.dp), + ) { + if (showScrollUp) { + GlassSurface( + shape = KaizenShapes.circle, + tintAlpha = GlassTiers.pill(isSystemInDarkTheme()), + shadowStack = KaizenShadows.level1, + cornerRadius = 18.dp, + modifier = Modifier + .size(36.dp) + .clickable { scope.launch { listState.animateScrollToItem(0) } }, + ) { + Box(Modifier.fillMaxSize(), contentAlignment = Alignment.Center) { + Icon(Icons.Rounded.KeyboardArrowUp, null, tint = MaterialTheme.colorScheme.onBackground, modifier = Modifier.size(20.dp)) + } + } + } + if (showScrollDown) { + GlassSurface( + shape = KaizenShapes.circle, + tintAlpha = GlassTiers.pill(isSystemInDarkTheme()), + shadowStack = KaizenShadows.level1, + cornerRadius = 18.dp, + modifier = Modifier + .size(36.dp) + .clickable { scope.launch { listState.animateScrollToItem(messages.size - 1) } }, + ) { + Box(Modifier.fillMaxSize(), contentAlignment = Alignment.Center) { + Icon(Icons.Rounded.KeyboardArrowDown, null, tint = MaterialTheme.colorScheme.onBackground, modifier = Modifier.size(20.dp)) + } + } + } + } + } } // 2. Floating Top Menu Button @@ -585,6 +638,10 @@ fun ChatScreen( onClick = { haptics.tick(); showModelSheet = true }, modifier = Modifier.weight(1f, fill = false), ) + if (usedTokens > 0 || messages.isNotEmpty()) { + Spacer(Modifier.width(8.dp)) + TokenBadge(used = usedTokens, contextLength = 1_000_000) + } } } diff --git a/app/src/main/java/dev/kaizen/app/net/StreamConsumer.kt b/app/src/main/java/dev/kaizen/app/net/StreamConsumer.kt index 8158851..f8d8c44 100644 --- a/app/src/main/java/dev/kaizen/app/net/StreamConsumer.kt +++ b/app/src/main/java/dev/kaizen/app/net/StreamConsumer.kt @@ -28,12 +28,20 @@ data class SearchSource( val snippet: String = "", ) +@Serializable +private data class UsageJson( + val i: Int = 0, + val o: Int = 0, +) + data class StreamState( val content: String = "", val reasoning: String = "", val tools: List = emptyList(), val query: String = "", val sources: List = emptyList(), + val inputTokens: Int = 0, + val outputTokens: Int = 0, ) object StreamConsumer { @@ -101,12 +109,15 @@ object StreamConsumer { private fun buildState(): StreamState { val queryStr = parseTrailingField(QUERY) val sourcesList = parseTrailingSources() + val usage = parseTrailingUsage() return StreamState( content = content.toString(), reasoning = reasoningBuf.toString(), tools = tools.toList(), query = queryStr, sources = sourcesList, + inputTokens = usage?.i ?: 0, + outputTokens = usage?.o ?: 0, ) } @@ -135,6 +146,15 @@ object StreamConsumer { } } + private fun parseTrailingUsage(): UsageJson? { + val full = trailingBuf.toString() + val idx = full.indexOf(USAGE) + if (idx == -1) return null + val json = full.substring(idx + 1).trim() + if (json.isEmpty()) return null + return try { lenientJson.decodeFromString(json) } catch (_: Exception) { null } + } + private fun parseToolEvents(buf: CharSequence, from: Int, to: Int) { val s = buf.toString() var pos = from