feat: load model prices from JSON files

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:

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")

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)
 

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")

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)