Skip to content

feat: add token usage and cost display to model responses #7

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sena-labs merged 1 commit into from
May 7, 2026
Merged

feat: add token usage and cost display to model responses #7

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
78 changes: 78 additions & 0 deletions openrouter_pipe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -120,6 +120,47 @@ def _format_citation_list(citations: Optional[List[str]]) -> str:
return ""


_CURRENCY_SYMBOLS = {
"USD": "$",
"EUR": "€",
"GBP": "£",
"JPY": "¥",
"CAD": "CA$",
"AUD": "A$",
}


def _format_cost_info(usage: dict, currency: str = "USD") -> str:
"""Format token usage and cost from an OpenRouter usage dict."""
if not usage:
return ""
prompt = usage.get("prompt_tokens", 0)
completion = usage.get("completion_tokens", 0)
total = usage.get("total_tokens", 0) or (prompt + completion)
cost = usage.get("cost")

token_str = f"{prompt:,} prompt + {completion:,} completion = {total:,} total"
parts = [f"**Tokens:** {token_str}"]

if cost is not None:
try:
cost_f = float(cost)
symbol = _CURRENCY_SYMBOLS.get(currency, f"{currency} ")
if cost_f == 0:
cost_str = f"{symbol}0.00"
elif cost_f < 0.0001:
cost_str = f"{symbol}{cost_f:.6f}"
elif cost_f < 0.01:
cost_str = f"{symbol}{cost_f:.5f}"
else:
cost_str = f"{symbol}{cost_f:.4f}"
parts.append(f"**Cost:** {cost_str}")
except (ValueError, TypeError):
pass

return f"\n\n---\n*{' · '.join(parts)}*"


class Pipe:
class Valves(BaseModel):
OPENROUTER_API_KEY: str = Field(
Expand Down Expand Up @@ -233,6 +274,27 @@ class Valves(BaseModel):
MAX_RETRIES: int = Field(
default=2, ge=0, description="Auto-retries on transient errors (with exponential backoff)"
)
SHOW_COST_INFO: bool = Field(
default=False,
description="Append token usage and cost to each response",
)
COST_CURRENCY: str = Field(
default=os.getenv("OPENROUTER_COST_CURRENCY", "USD"),
description="Currency label shown in cost display (display only; OpenRouter bills in USD)",
json_schema_extra={
"input": {
"type": "select",
"options": [
{"value": "USD", "label": "USD ($)"},
{"value": "EUR", "label": "EUR (€)"},
{"value": "GBP", "label": "GBP (£)"},
{"value": "JPY", "label": "JPY (¥)"},
{"value": "CAD", "label": "CAD (CA$)"},
{"value": "AUD", "label": "AUD (A$)"},
],
}
},
)

@field_validator("OPENROUTER_BASE_URL")
@classmethod
Expand Down Expand Up @@ -768,6 +830,12 @@ def _non_stream_response(self, headers: dict, payload: dict) -> str:

if rendered_citations:
final_parts.append(rendered_citations)

if self.valves.SHOW_COST_INFO:
cost_info = _format_cost_info(res.get("usage", {}), self.valves.COST_CURRENCY)
if cost_info:
final_parts.append(cost_info)

return "".join(final_parts)
except requests.exceptions.Timeout:
return f"OpenRouter Error: Request timed out after {self.valves.REQUEST_TIMEOUT}s. Try increasing REQUEST_TIMEOUT or retry."
Expand All @@ -785,6 +853,7 @@ def _stream_response(
response = None
in_think = False
latest_citations: List[str] = []
latest_usage: dict = {}

def _close_think_tag():
nonlocal in_think
Expand Down Expand Up @@ -820,6 +889,10 @@ def _close_think_tag():
yield f"\n\nOpenRouter Error: {msg}"
return

usage_data = chunk.get("usage")
if usage_data:
latest_usage = usage_data

citations = chunk.get("citations")
if citations is not None:
latest_citations = citations
Expand Down Expand Up @@ -849,6 +922,11 @@ def _close_think_tag():
rendered_citations = _format_citation_list(latest_citations)
if rendered_citations:
yield rendered_citations

if self.valves.SHOW_COST_INFO:
cost_info = _format_cost_info(latest_usage, self.valves.COST_CURRENCY)
if cost_info:
yield cost_info
except requests.exceptions.Timeout:
close_tag = _close_think_tag()
if close_tag:
Expand Down
180 changes: 180 additions & 0 deletions test_pipe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -145,6 +145,8 @@ def _section(title: str):
_assert(v.ENABLE_CACHE_CONTROL is False, "ENABLE_CACHE_CONTROL false")
_assert(v.REQUEST_TIMEOUT == 90, "REQUEST_TIMEOUT 90")
_assert(v.MAX_RETRIES == 2, "MAX_RETRIES 2")
_assert(v.SHOW_COST_INFO is False, "SHOW_COST_INFO false by default")
_assert(v.COST_CURRENCY == "USD", "COST_CURRENCY USD by default")

try:
Pipe.Valves(REQUEST_TIMEOUT=-1)
Expand Down Expand Up @@ -1854,6 +1856,184 @@ async def _test_pipe_no_msgs_key():
"retryable: stream=False forwarded to requests.Session.post",
)

# ── 33. _format_cost_info() and SHOW_COST_INFO ──────────────────────────────

_section("33. _format_cost_info() and SHOW_COST_INFO")

_format_cost_info = mod._format_cost_info
_CURRENCY_SYMBOLS = mod._CURRENCY_SYMBOLS

# 33a. Empty usage dict → empty string
_assert(_format_cost_info({}) == "", "_format_cost_info: empty dict → empty")
_assert(_format_cost_info({}, "EUR") == "", "_format_cost_info: empty dict with currency → empty")

# 33b. Usage with tokens only (no cost) → tokens shown
_result_tokens_only = _format_cost_info({"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150})
_assert("100" in _result_tokens_only, "_format_cost_info: prompt tokens present")
_assert("50" in _result_tokens_only, "_format_cost_info: completion tokens present")
_assert("150" in _result_tokens_only, "_format_cost_info: total tokens present")
_assert("Tokens" in _result_tokens_only, "_format_cost_info: Tokens label present")
_assert("Cost" not in _result_tokens_only, "_format_cost_info: no cost when field absent")

# 33c. Zero cost → $0.00
_result_free = _format_cost_info({"prompt_tokens": 200, "completion_tokens": 100, "total_tokens": 300, "cost": 0})
_assert("$0.00" in _result_free, "_format_cost_info: zero cost → $0.00")
_assert("Cost" in _result_free, "_format_cost_info: Cost label present for zero cost")

# 33d. Micro cost < 0.0001 → 6 decimal places
_result_micro = _format_cost_info({"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15, "cost": 0.000005})
_assert("$0.000005" in _result_micro, "_format_cost_info: micro cost 6 decimal places")

# 33e. Small cost 0.0001–0.01 → 5 decimal places
_result_small = _format_cost_info({"prompt_tokens": 50, "completion_tokens": 25, "total_tokens": 75, "cost": 0.001234})
_assert("$0.00123" in _result_small, "_format_cost_info: small cost 5 decimal places")

# 33f. Normal cost >= 0.01 → 4 decimal places
_result_normal = _format_cost_info({"prompt_tokens": 500, "completion_tokens": 200, "total_tokens": 700, "cost": 0.05670})
_assert("$0.0567" in _result_normal, "_format_cost_info: normal cost 4 decimal places")

# 33g. EUR currency symbol
_result_eur = _format_cost_info({"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150, "cost": 0.01}, "EUR")
_assert("€" in _result_eur, "_format_cost_info: EUR symbol shown")

# 33h. GBP currency symbol
_result_gbp = _format_cost_info({"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150, "cost": 0.01}, "GBP")
_assert("£" in _result_gbp, "_format_cost_info: GBP symbol shown")

# 33i. Unknown currency → uses currency string as prefix
_result_unknown = _format_cost_info({"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15, "cost": 0.01}, "XYZ")
_assert("XYZ " in _result_unknown, "_format_cost_info: unknown currency → code as prefix")

# 33j. Invalid cost value (string) → tokens shown, cost silently skipped
_result_bad_cost = _format_cost_info({"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15, "cost": "invalid"})
_assert("Tokens" in _result_bad_cost, "_format_cost_info: invalid cost → tokens still shown")
_assert("Cost" not in _result_bad_cost, "_format_cost_info: invalid cost → cost silently skipped")

# 33k. total_tokens missing → computed as prompt + completion
_result_no_total = _format_cost_info({"prompt_tokens": 80, "completion_tokens": 20, "cost": 0.005})
_assert("100" in _result_no_total, "_format_cost_info: total computed from prompt+completion when missing")

# 33l. Output format: starts with separator, italic, bold labels
_result_fmt = _format_cost_info({"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2, "cost": 0.001})
_assert(_result_fmt.startswith("\n\n---\n"), "_format_cost_info: starts with separator")
_assert("*" in _result_fmt, "_format_cost_info: uses italic (asterisks)")
_assert("**Tokens:**" in _result_fmt, "_format_cost_info: bold Tokens label")
_assert("**Cost:**" in _result_fmt, "_format_cost_info: bold Cost label")

# 33m. Large numbers use comma separators
_result_large = _format_cost_info({"prompt_tokens": 1250, "completion_tokens": 342, "total_tokens": 1592, "cost": 0.00234})
_assert("1,250" in _result_large, "_format_cost_info: large prompt token count formatted with comma")
_assert("1,592" in _result_large, "_format_cost_info: large total formatted with comma")

# 33n. CURRENCY_SYMBOLS dict completeness
for _sym_key in ("USD", "EUR", "GBP", "JPY", "CAD", "AUD"):
_assert(_sym_key in _CURRENCY_SYMBOLS, f"_CURRENCY_SYMBOLS: {_sym_key} present")

# ── 33 (cont). Non-stream SHOW_COST_INFO integration ────────────────────────

pipe = Pipe()
pipe.valves = Pipe.Valves(OPENROUTER_API_KEY="k")

# 33o. SHOW_COST_INFO=False (default) → no cost appended
_mock_cost_off = MagicMock()
_mock_cost_off.json.return_value = {
"choices": [{"message": {"content": "Hello"}}],
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15, "cost": 0.001},
}
with patch.object(pipe, "_retryable_request", return_value=_mock_cost_off):
_result_off = pipe._non_stream_response({}, {})
_assert("---" not in _result_off, "non-stream SHOW_COST_INFO=False: no separator appended")
_assert("Tokens" not in _result_off, "non-stream SHOW_COST_INFO=False: no Tokens line")

# 33p. SHOW_COST_INFO=True → cost appended after content
pipe.valves = Pipe.Valves(OPENROUTER_API_KEY="k", SHOW_COST_INFO=True, COST_CURRENCY="USD")
_mock_cost_on = MagicMock()
_mock_cost_on.json.return_value = {
"choices": [{"message": {"content": "Hello"}}],
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15, "cost": 0.00123},
}
with patch.object(pipe, "_retryable_request", return_value=_mock_cost_on):
_result_on = pipe._non_stream_response({}, {})
_assert("Hello" in _result_on, "non-stream SHOW_COST_INFO=True: content preserved")
_assert("Tokens" in _result_on, "non-stream SHOW_COST_INFO=True: Tokens label present")
_assert("$" in _result_on, "non-stream SHOW_COST_INFO=True: cost shown in USD")
_assert("10" in _result_on, "non-stream SHOW_COST_INFO=True: prompt token count present")

# 33q. SHOW_COST_INFO=True with EUR currency
pipe.valves = Pipe.Valves(OPENROUTER_API_KEY="k", SHOW_COST_INFO=True, COST_CURRENCY="EUR")
_mock_cost_eur = MagicMock()
_mock_cost_eur.json.return_value = {
"choices": [{"message": {"content": "Ciao"}}],
"usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8, "cost": 0.05},
}
with patch.object(pipe, "_retryable_request", return_value=_mock_cost_eur):
_result_eur_resp = pipe._non_stream_response({}, {})
_assert("€" in _result_eur_resp, "non-stream SHOW_COST_INFO=True EUR: euro symbol shown")

# 33r. SHOW_COST_INFO=True but response has no usage → no cost appended (no crash)
pipe.valves = Pipe.Valves(OPENROUTER_API_KEY="k", SHOW_COST_INFO=True)
_mock_no_usage = MagicMock()
_mock_no_usage.json.return_value = {
"choices": [{"message": {"content": "No usage data"}}],
}
with patch.object(pipe, "_retryable_request", return_value=_mock_no_usage):
_result_no_usage = pipe._non_stream_response({}, {})
_assert("No usage data" in _result_no_usage, "non-stream no usage: content preserved")
_assert("Tokens" not in _result_no_usage, "non-stream no usage: no Tokens line (no crash)")

# ── 33 (cont). Stream SHOW_COST_INFO integration ─────────────────────────────

pipe = Pipe()
pipe.valves = Pipe.Valves(OPENROUTER_API_KEY="k", SHOW_COST_INFO=True, COST_CURRENCY="USD")

# 33s. Usage in final SSE chunk → cost appended after stream
_usage_chunk = {"prompt_tokens": 150, "completion_tokens": 75, "total_tokens": 225, "cost": 0.00567}
sse_with_usage = [
b"data: " + json.dumps({"choices": [{"delta": {"content": "Answer"}}]}).encode(),
b"data: " + json.dumps({"choices": [{"delta": {}}], "usage": _usage_chunk}).encode(),
b"data: [DONE]",
]
with patch.object(pipe, "_retryable_request", return_value=_make_sse_response(sse_with_usage)):
_stream_output = list(pipe._stream_response({}, {}))
_stream_full = "".join(_stream_output)
_assert("Answer" in _stream_full, "stream SHOW_COST_INFO=True: content preserved")
_assert("Tokens" in _stream_full, "stream SHOW_COST_INFO=True: Tokens label in output")
_assert("150" in _stream_full, "stream SHOW_COST_INFO=True: prompt token count present")
_assert("$" in _stream_full, "stream SHOW_COST_INFO=True: cost in USD")

# 33t. SHOW_COST_INFO=False → usage in chunk but no cost appended
pipe.valves = Pipe.Valves(OPENROUTER_API_KEY="k", SHOW_COST_INFO=False)
sse_no_cost_display = [
b"data: " + json.dumps({"choices": [{"delta": {"content": "Reply"}}]}).encode(),
b"data: " + json.dumps({"choices": [{"delta": {}}], "usage": _usage_chunk}).encode(),
b"data: [DONE]",
]
with patch.object(pipe, "_retryable_request", return_value=_make_sse_response(sse_no_cost_display)):
_stream_off_output = list(pipe._stream_response({}, {}))
_stream_off_full = "".join(_stream_off_output)
_assert("Reply" in _stream_off_full, "stream SHOW_COST_INFO=False: content preserved")
_assert("Tokens" not in _stream_off_full, "stream SHOW_COST_INFO=False: no Tokens line")

# 33u. Stream with no usage chunk → no cost appended, no crash
pipe.valves = Pipe.Valves(OPENROUTER_API_KEY="k", SHOW_COST_INFO=True)
sse_no_usage_chunk = [
b"data: " + json.dumps({"choices": [{"delta": {"content": "Text"}}]}).encode(),
b"data: [DONE]",
]
with patch.object(pipe, "_retryable_request", return_value=_make_sse_response(sse_no_usage_chunk)):
_stream_nu_output = list(pipe._stream_response({}, {}))
_stream_nu_full = "".join(_stream_nu_output)
_assert("Text" in _stream_nu_full, "stream no usage chunk: content preserved")
_assert("Tokens" not in _stream_nu_full, "stream no usage chunk: no cost line (no crash)")

# 33v. COST_CURRENCY valve uses select with 6 options
_cc_field = Pipe.Valves.model_fields["COST_CURRENCY"]
_cc_options = _cc_field.json_schema_extra.get("input", {}).get("options", [])
_assert(len(_cc_options) == 6, "COST_CURRENCY: 6 currency options")
_cc_values = [o["value"] for o in _cc_options]
_assert("USD" in _cc_values, "COST_CURRENCY options: USD present")
_assert("EUR" in _cc_values, "COST_CURRENCY options: EUR present")

# ══════════════════════════════════════════════════════════════════════════════
# Summary
# ══════════════════════════════════════════════════════════════════════════════
Expand Down
Loading