Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 20 additions & 9 deletions src/ucode/usage.py
Original file line number Diff line number Diff line change
Expand Up @@ -636,24 +636,32 @@ def render_usage_summary(
if usage_day == today:
daily_total += token_total

today_text = format_token_count(daily_total)
weekly_text = format_token_count(weekly_total)
monthly_text = format_token_count(monthly_total)
token_table = render_box_table(
["Today", "Last 7 days", "Last 30 days"],
[[f"{today_text} tokens", f"{weekly_text} tokens", f"{monthly_text} tokens"]],
)
for token_text in {today_text, weekly_text, monthly_text}:
plain_value = f"{token_text} tokens"
token_table = token_table.replace(plain_value, value(plain_value))
lines = [
heading(f"Usage Summary for {requester_name}"),
"",
"[bold green]✓[/bold green] Databricks AI Gateway usage",
f"{label('Today:')} {value(format_token_count(daily_total) + ' tokens')}",
f"{label('Last 7 days:')} {value(format_token_count(weekly_total) + ' tokens')}",
f"{label('Last 30 days:')} {value(format_token_count(monthly_total) + ' tokens')}",
"[bold green]✓[/bold green] [bold]Databricks AI Gateway Token Usage[/bold]",
token_table,
]
if active_tools_last_week:
tool_text = ", ".join(tool_displays[tool] for tool in active_tools_last_week)
lines.append(f"{label('Active tools:')} {value(tool_text)}")
tool_text = " · ".join(tool_displays[tool] for tool in active_tools_last_week)
lines.extend(["", f" 🤖 {label('Top Harnesses:')} {value(tool_text)}"])
if weekly_model_usage:
top_models = sorted(
weekly_model_usage.values(),
key=lambda u: (-u.total, u.name.lower()),
)[:3]
models_text = ", ".join(usage.name for usage in top_models)
lines.append(f"{label('Top models this week:')} {value(models_text)}")
models_text = " · ".join(usage.name for usage in top_models)
lines.append(f"{label('Top Models:')} {value(models_text)}")
weekly_cost = sum(
(
model_usage_cost(usage, price_lookup) or Decimal(0)
Expand All @@ -662,7 +670,10 @@ def render_usage_summary(
Decimal(0),
)
if weekly_cost > 0:
lines.append(f"{label('Est. cost (7 days):')} {value(format_cost_usd(weekly_cost))}")
lines.append(
f" 💰 {label('Estimated Cost for Last 7 Days:')} "
f"{value(format_cost_usd(weekly_cost))}"
)
lines.extend(render_budget_lines(budget_spend))
return "\n".join(lines)

Expand Down
48 changes: 39 additions & 9 deletions tests/test_usage.py
Original file line number Diff line number Diff line change
Expand Up @@ -513,6 +513,11 @@ def test_today_total(self):
records = [self._make_record(0, "claude", 5000)]
result = render_usage_summary(records, "user", {"claude": "Claude Code"})
assert "5.0K" in result
assert "Databricks AI Gateway Token Usage" in result
assert "┃ Today" in result
assert "Last 7 days" in result
assert "Last 30 days" in result
assert "│ [cyan]5.0K tokens[/cyan]" in result

def test_weekly_total_includes_past_week(self):
records = [
Expand All @@ -525,9 +530,12 @@ def test_weekly_total_includes_past_week(self):
assert "3.0K" in result or "3" in result

def test_active_tools_listed(self):
records = [self._make_record(0, "claude", 1000)]
result = render_usage_summary(records, "user", {"claude": "Claude Code"})
assert "Claude Code" in result
records = [
self._make_record(0, "claude", 1000),
self._make_record(0, "codex", 1000),
]
result = render_usage_summary(records, "user", {"claude": "Claude Code", "codex": "Codex"})
assert "🤖 [bold]Top Harnesses:[/bold] [cyan]Claude Code · Codex[/cyan]" in result

def test_top_models_listed(self):
records = [
Expand All @@ -541,6 +549,27 @@ def test_top_models_listed(self):
result = render_usage_summary(records, "user", {"claude": "Claude Code"})
assert "claude-sonnet-4" in result

def test_estimated_cost_label(self):
records = [
{
"tool": "claude",
"usage_day": date.today(),
"total_tokens_used": 1_000_000,
"model_tokens": (
'[{"model":"databricks-claude-sonnet-4","tokens":1000000,'
'"input":1000000,"output":0}]'
),
}
]
result = render_usage_summary(
records,
"user",
{"claude": "Claude Code"},
price_lookup={"claudesonnet4": ModelPrice(Decimal("4.12"), None, None)},
)
assert "Estimated Cost for Last 7 Days:" in result
assert "$4.12" in result

def test_includes_budget_spend_when_available(self):
records = [self._make_record(0, "claude", 1000)]
result = render_usage_summary(
Expand Down Expand Up @@ -583,8 +612,8 @@ def test_top_models_uses_per_model_token_totals(self):
)
# Top-models line is full names only, ranked by per-model token totals
# (claude-opus-4 236.1K > gpt-5 13.3K > claude-haiku-4.5 920).
assert "Top models this week:" in result
assert "claude-opus-4, gpt-5, claude-haiku-4.5" in result
assert "⭐ [bold]Top Models:[/bold]" in result
assert "claude-opus-4 · gpt-5 · claude-haiku-4.5" in result
# No token counts in this line — those live in the per-model table.
assert "236.1K" not in result

Expand Down Expand Up @@ -680,11 +709,12 @@ def fake_render_box_table(headers, table_rows, max_widths=None):
"Using SQL warehouse `wh` (RUNNING).",
f"No usage for Claude Code in the last {USAGE_BREAKDOWN_DAYS} days.",
]
assert len(rendered_tables) == 1
assert len(rendered_tables) == 2
assert rendered_tables[0] == [["100 tokens", "100 tokens", "300 tokens"]]
# One per-model row for codex: full model name "gpt-5", 1 request, 80 input / 20 output.
assert rendered_tables[0][0][0] == "gpt-5"
assert rendered_tables[0][0][1] == "1"
assert rendered_tables[0][0][2] == "80"
assert rendered_tables[1][0][0] == "gpt-5"
assert rendered_tables[1][0][1] == "1"
assert rendered_tables[1][0][2] == "80"
assert "gemini" not in "\n".join(printed).lower()
assert "900" not in "\n".join(printed)

Expand Down
Loading