diff --git a/README.md b/README.md index 39bc3ef..dbca854 100644 --- a/README.md +++ b/README.md @@ -37,9 +37,31 @@ $ tokenmeter count system.txt -m gpt-4o # one file $ tokenmeter count prompts/ -m gpt-4o-mini # every text file in a directory $ cat prompt.txt | tokenmeter count - -m gpt-4o # standard input $ tokenmeter count p.txt --output-tokens 500 # include an assumed completion +$ tokenmeter count p.txt --prices-file prices.json # add or override model prices $ tokenmeter models # list models and prices ``` +### Custom price files + +Use `--prices-file` with `count`, `budget`, or `models` to add new model prices or +override the built-in table. Set `TOKENMETER_PRICES_FILE` to use the same custom +price file for every command without repeating the flag; an explicit +`--prices-file` value takes precedence over the environment variable. The file is +JSON keyed by model name: + +```json +{ + "local-model": { + "encoding": "cl100k_base", + "input_per_mtok": 1.0, + "output_per_mtok": 3.0 + } +} +``` + +Prices are in USD per 1,000,000 tokens. `encoding` must be a tiktoken encoding +name such as `cl100k_base` or `o200k_base`. + ### As a budget gate Fail a build when a prompt set would cost more than you allow: diff --git a/src/tokenmeter/cli.py b/src/tokenmeter/cli.py index 75c7742..7045b03 100644 --- a/src/tokenmeter/cli.py +++ b/src/tokenmeter/cli.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +import os import sys import typer @@ -15,8 +16,10 @@ from tokenmeter.meter import measure, over_budget, total_cost from tokenmeter.pricing import ( PRICES_AS_OF, + PriceFileError, UnknownModel, known_models, + load_prices_file, price_for, ) from tokenmeter.render import measurements_to_json, render_table @@ -53,6 +56,17 @@ def main( """tokenmeter command-line interface.""" +def _load_prices_or_exit(prices_file: str | None) -> None: + prices_path = prices_file or os.environ.get("TOKENMETER_PRICES_FILE") + if prices_path is None: + return + try: + load_prices_file(prices_path) + except PriceFileError as exc: + _err.print(f"tokenmeter: {exc}") + raise typer.Exit(EXIT_BAD_INPUT) from exc + + def _collect(paths, model, output_tokens): inputs = read_inputs(paths) encoder = encoder_for_model(model) @@ -69,9 +83,15 @@ def count( 0, "--output-tokens", help="Assumed completion tokens, for cost only." ), as_json: bool = typer.Option(False, "--json", help="Emit JSON."), + prices_file: str | None = typer.Option( + None, + "--prices-file", + help="JSON file with model prices to add or override.", + ), ) -> None: """Count tokens and estimate cost for one or more inputs.""" + _load_prices_or_exit(prices_file) try: measurements = _collect(paths, model, output_tokens) except UnknownModel as exc: @@ -93,9 +113,15 @@ def budget( max_cost: float = typer.Option(..., "--max-cost", help="Fail above this USD cost."), model: str = typer.Option("gpt-4o", "--model", "-m", help="Model to price against."), output_tokens: int = typer.Option(0, "--output-tokens", help="Assumed completion tokens."), + prices_file: str | None = typer.Option( + None, + "--prices-file", + help="JSON file with model prices to add or override.", + ), ) -> None: """Fail when the estimated cost of the inputs exceeds a budget.""" + _load_prices_or_exit(prices_file) try: measurements = _collect(paths, model, output_tokens) except UnknownModel as exc: @@ -112,9 +138,16 @@ def budget( @app.command("models") -def models() -> None: +def models( + prices_file: str | None = typer.Option( + None, + "--prices-file", + help="JSON file with model prices to add or override.", + ), +) -> None: """List the known models and their prices.""" + _load_prices_or_exit(prices_file) title = f"prices as of {PRICES_AS_OF} (USD per 1M tokens)" table = Table(box=None, pad_edge=False, title=title) table.add_column("model") diff --git a/src/tokenmeter/pricing.py b/src/tokenmeter/pricing.py index 90f66f2..8fd656d 100644 --- a/src/tokenmeter/pricing.py +++ b/src/tokenmeter/pricing.py @@ -7,7 +7,10 @@ from __future__ import annotations +import json from dataclasses import dataclass +from pathlib import Path +from typing import Any PRICES_AS_OF = "2025-08-01" @@ -44,6 +47,65 @@ def __str__(self) -> str: return f"unknown model: {self.model}" +class PriceFileError(ValueError): + """Raised when an external price file cannot be loaded.""" + + +def _as_number(value: Any, *, model: str, field: str) -> float: + if isinstance(value, bool) or not isinstance(value, int | float): + raise PriceFileError(f"model {model!r} field {field!r} must be a number") + return float(value) + + +def _as_string(value: Any, *, model: str, field: str) -> str: + if not isinstance(value, str) or not value: + raise PriceFileError(f"model {model!r} field {field!r} must be a non-empty string") + return value + + +def load_prices_file(path: str | Path) -> None: + """Merge model prices from a JSON file into the active price table. + + The JSON file must contain a top-level object keyed by model name. Each value + must include ``encoding``, ``input_per_mtok`` and ``output_per_mtok``. + Existing model names are replaced, and new model names are added. + """ + + price_path = Path(path) + try: + raw = json.loads(price_path.read_text()) + except OSError as exc: + raise PriceFileError(f"could not read price file {price_path}: {exc}") from exc + except json.JSONDecodeError as exc: + raise PriceFileError(f"could not parse price file {price_path}: {exc}") from exc + + if not isinstance(raw, dict): + raise PriceFileError("price file must contain a JSON object keyed by model name") + + loaded: dict[str, ModelPrice] = {} + for model, payload in raw.items(): + if not isinstance(model, str) or not model: + raise PriceFileError("price file model names must be non-empty strings") + if not isinstance(payload, dict): + raise PriceFileError(f"model {model!r} must map to a JSON object") + missing = {"encoding", "input_per_mtok", "output_per_mtok"} - payload.keys() + if missing: + fields = ", ".join(sorted(missing)) + raise PriceFileError(f"model {model!r} is missing required field(s): {fields}") + loaded[model] = ModelPrice( + model=model, + encoding=_as_string(payload["encoding"], model=model, field="encoding"), + input_per_mtok=_as_number( + payload["input_per_mtok"], model=model, field="input_per_mtok" + ), + output_per_mtok=_as_number( + payload["output_per_mtok"], model=model, field="output_per_mtok" + ), + ) + + _PRICES.update(loaded) + + def known_models() -> list[str]: return sorted(_PRICES) diff --git a/tests/test_cli.py b/tests/test_cli.py index abecb92..550893e 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -42,6 +42,101 @@ def test_count_unknown_model_is_bad_input(tmp_path): assert result.exit_code == cli_module.EXIT_BAD_INPUT +def test_count_uses_custom_prices_file(tmp_path): + f = tmp_path / "p.txt" + f.write_text("one two three four") + prices = tmp_path / "prices.json" + prices.write_text( + '{"local-model": {"encoding": "cl100k_base", "input_per_mtok": 1.0, ' + '"output_per_mtok": 3.0}}' + ) + + result = runner.invoke( + cli_module.app, + [ + "count", + str(f), + "--model", + "local-model", + "--output-tokens", + "2", + "--prices-file", + str(prices), + "--json", + ], + ) + + assert result.exit_code == 0 + payload = json.loads(result.stdout) + assert payload["inputs"][0]["model"] == "local-model" + assert payload["inputs"][0]["input_cost"] == 0.000004 + assert payload["inputs"][0]["output_cost"] == 0.000006 + assert payload["total_cost"] == 0.00001 + + +def test_count_uses_prices_file_from_environment(tmp_path, monkeypatch): + f = tmp_path / "p.txt" + f.write_text("one two") + prices = tmp_path / "prices.json" + prices.write_text( + '{"env-model": {"encoding": "cl100k_base", "input_per_mtok": 2.0, ' + '"output_per_mtok": 4.0}}' + ) + monkeypatch.setenv("TOKENMETER_PRICES_FILE", str(prices)) + + result = runner.invoke(cli_module.app, ["count", str(f), "--model", "env-model", "--json"]) + + assert result.exit_code == 0 + payload = json.loads(result.stdout) + assert payload["inputs"][0]["model"] == "env-model" + assert payload["inputs"][0]["input_cost"] == 0.000004 + + +def test_prices_file_option_takes_precedence_over_environment(tmp_path, monkeypatch): + f = tmp_path / "p.txt" + f.write_text("one two") + env_prices = tmp_path / "env-prices.json" + env_prices.write_text('{"env-only": {"encoding": "cl100k_base"}}') + option_prices = tmp_path / "option-prices.json" + option_prices.write_text( + '{"option-model": {"encoding": "cl100k_base", "input_per_mtok": 1.5, ' + '"output_per_mtok": 2.5}}' + ) + monkeypatch.setenv("TOKENMETER_PRICES_FILE", str(env_prices)) + + result = runner.invoke( + cli_module.app, + [ + "count", + str(f), + "--model", + "option-model", + "--prices-file", + str(option_prices), + "--json", + ], + ) + + assert result.exit_code == 0 + payload = json.loads(result.stdout) + assert payload["inputs"][0]["model"] == "option-model" + + +def test_prices_file_schema_error_is_bad_input(tmp_path): + f = tmp_path / "p.txt" + f.write_text("hello") + prices = tmp_path / "prices.json" + prices.write_text('{"broken": {"encoding": "cl100k_base"}}') + + result = runner.invoke( + cli_module.app, + ["count", str(f), "--model", "broken", "--prices-file", str(prices)], + ) + + assert result.exit_code == cli_module.EXIT_BAD_INPUT + assert "missing required" in result.stderr + + def test_budget_passes_under_limit(tmp_path): f = tmp_path / "p.txt" f.write_text("one two three") diff --git a/tests/test_pricing.py b/tests/test_pricing.py index ea80b9a..24732a1 100644 --- a/tests/test_pricing.py +++ b/tests/test_pricing.py @@ -3,9 +3,11 @@ import pytest from tokenmeter.pricing import ( + PriceFileError, UnknownModel, input_cost, known_models, + load_prices_file, output_cost, price_for, total_cost, @@ -41,3 +43,27 @@ def test_total_cost_adds_input_and_output(): def test_embeddings_have_no_output_cost(): assert output_cost("text-embedding-3-small", 1_000_000) == 0.0 + + +def test_load_prices_file_adds_custom_model(tmp_path): + prices = tmp_path / "prices.json" + prices.write_text( + '{"custom-fast": {"encoding": "cl100k_base", "input_per_mtok": 1.25, ' + '"output_per_mtok": 2.5}}' + ) + + load_prices_file(prices) + + assert "custom-fast" in known_models() + price = price_for("custom-fast") + assert price.encoding == "cl100k_base" + assert input_cost("custom-fast", 1_000_000) == pytest.approx(1.25) + assert output_cost("custom-fast", 1_000_000) == pytest.approx(2.5) + + +def test_load_prices_file_rejects_bad_schema(tmp_path): + prices = tmp_path / "prices.json" + prices.write_text('{"broken": {"encoding": "cl100k_base"}}') + + with pytest.raises(PriceFileError, match="missing required"): + load_prices_file(prices)