diff --git a/src/storied/cli.py b/src/storied/cli.py index d305d72..cf93927 100644 --- a/src/storied/cli.py +++ b/src/storied/cli.py @@ -381,57 +381,90 @@ def cmd_play(args: argparse.Namespace) -> int: stats = engine.get_context_stats() console.print() - # Header with model, time, and context estimate - context_k = stats["context_total"] / 1000 + # Header: real API usage if available, estimate otherwise limit_k = stats["model_limit"] / 1000 game_time = engine.get_current_time() + if stats["last_input"] > 0: + input_k = stats["last_input"] / 1000 + usage_str = f"{input_k:.1f}k/{limit_k:.0f}k tokens (last turn)" + else: + context_k = stats["context_total"] / 1000 + usage_str = f"~{context_k:.1f}k/{limit_k:.0f}k system tokens (estimated)" console.print( - f"[bold]Context Usage[/bold] [dim]({game_time})[/dim] " - f"[dim]{engine.model} · ~{context_k:.1f}k/{limit_k:.0f}k system tokens[/dim]" + f"[bold]Context Usage[/bold] [dim]({game_time})[/dim]" + f" [dim]{engine.model} · {usage_str}[/dim]" ) - # Build visual bar (30 chars wide) - bar_width = 30 - limit = stats["model_limit"] - bar_parts = [ + # Color mapping for context sections + _SECTION_COLORS: dict[str, str] = { + "Style": "dim", + "Character": "green", + "Log": "bright_cyan", + "Transcript": "blue", + "Session": "yellow", + "PlayerKnowledge": "bright_green", + "Location": "cyan", + "Initiative": "red", + } + + # Build section list from all context parts + sections: list[tuple[str, int, str]] = [ ("DM Instructions", stats["system_prompt"], "bright_blue"), - ("Character", stats["context_parts"].get("Character", 0), "green"), - ("Log", stats["context_parts"].get("Log", 0), "bright_cyan"), - ("Session", stats["context_parts"].get("Session", 0), "yellow"), ] - if "Location" in stats["context_parts"]: - bar_parts.append(("Location", stats["context_parts"]["Location"], "cyan")) - for key, val in stats["context_parts"].items(): + for key, tokens in stats["context_parts"].items(): if key.startswith("Entity:") or key.startswith("Linked:"): - bar_parts.append((key.split(":", 1)[1], val, "magenta")) - - bar = "" - legend_items = [] - for name, tokens, color in bar_parts: - if tokens > 0: - width = max(1, int((tokens / limit) * bar_width)) - bar += f"[{color}]{'█' * width}[/{color}]" - pct_part = (tokens / limit) * 100 - legend_items.append((name, tokens, pct_part, color)) - - used_width = sum(max(1, int((t / limit) * bar_width)) for _, t, _ in bar_parts if t > 0) - remaining_width = bar_width - used_width - if remaining_width > 0: - bar += f"[dim]{'░' * remaining_width}[/dim]" - - console.print(bar) - - for name, tokens, pct_part, color in legend_items: - tokens_str = f"{tokens:,}" if tokens < 1000 else f"{tokens/1000:.1f}k" - console.print(f" [{color}]█[/{color}] {name}: [dim]{tokens_str} ({pct_part:.1f}%)[/dim]") - - # Session totals + label = key.split(":", 1)[1] + color = "magenta" if key.startswith("Entity:") else "dark_magenta" + else: + label = key + color = _SECTION_COLORS.get(key, "white") + sections.append((label, tokens, color)) + + # 40x5 grid (200 cells = 200k tokens, 1 cell = 1k) + grid_w, grid_h = 40, 5 + total_cells = grid_w * grid_h + cells: list[str] = [] + legend_items: list[tuple[str, int, str]] = [] + for name, tokens, color in sections: + n_cells = round(tokens / 1_000) + if tokens > 0 and n_cells == 0: + n_cells = 1 + if n_cells > 0: + cells.extend([color] * n_cells) + legend_items.append((name, tokens, color)) + + # Pad remaining cells + remaining = total_cells - len(cells) + if remaining > 0: + cells.extend(["dim"] * remaining) + + # Render grid + for row in range(grid_h): + line = "" + for col in range(grid_w): + idx = row * grid_w + col + c = cells[idx] if idx < len(cells) else "dim" + char = "█" if c != "dim" else "░" + line += f"[{c}]{char}[/{c}]" + console.print(line) + + # Legend + for name, tokens, color in legend_items: + tokens_str = f"~{tokens:,}" if tokens < 1_000 else f"~{tokens / 1_000:.1f}k" + console.print(f" [{color}]█[/{color}] {name}: [dim]{tokens_str}[/dim]") + + # Session totals (actual API counts) if stats["total_input"] > 0: console.print() - console.print("[dim]Session totals:[/dim]") + last_in = f"{stats['last_input']:,}" + last_out = f"{stats['last_output']:,}" + total_in = f"{stats['total_input']:,}" + total_out = f"{stats['total_output']:,}" + console.print( + f" [dim]Last turn: {last_in} in · {last_out} out[/dim]" + ) console.print( - f" [dim]Tokens: {stats['total_input']:,} in · " - f"{stats['total_output']:,} out[/dim]" + f" [dim]Session: {total_in} in · {total_out} out[/dim]" ) console.print()