Skip to content

[Fix] Let decision backends declare billable input tokens - #17

Open
alexwong10 wants to merge 2 commits into
ThinkFlowLab:mainfrom
alexwong10:fix/rail-local-model-costs
Open

alexwong10 wants to merge 2 commits into
ThinkFlowLab:mainfrom
alexwong10:fix/rail-local-model-costs

Conversation

@alexwong10

@alexwong10 alexwong10 commented Sep 26, 2026 •

Copy link
Copy Markdown
Contributor

Why

A rail run with --model laya reported 300 local input tokens as $0.000013 of Jev API charges. Checking model.name == "jev" to decide billing would also miss a future paid backend with a different name.

How

DecisionModel.bills_input_tokens is a required backend declaration. Jev opts in; Laya, Cua, Random, Rule, and Scripted opt out. Rail, browser, and tool reports read the flag, and ToolDecisionModel carries it through its wrapper. Tool accounting excludes llm fallback ticks, which are already priced as chat tokens. A true flag uses the existing JEV_USD_PER_INPUT_TOKEN rate.

What

Per-decision token usage remains recorded for every backend. Free backends report zero API input charges, while a backend that opts in is billed across all three fronts regardless of its name. Regression tests cover the Jev and Laya adapters, a billed backend named scripted in each front, and an explicit billing declaration on every built-in backend.

Verification

  • Windows / Python 3.11.16: pytest -q — 495 passed, 41 skipped, 80 subtests passed.
  • Passed: ruff format --check ., ruff check ., ty check, uv lock --check --offline, uv build --offline, and scripts/smoke.sh in Git Bash.

@Yunaik Yunaik left a comment •

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

really good find! Thank you.

Please make the solution more long-lasting instead of string matching

rails.py:124 decides pricing with decision_model.name == "jev". Every new backend with a price would need another string check there. Ask the model instead. DecisionModel already declares capability flags (supports_images, deterministic). Add a billing flag next to them, with no default so that every backend has to state it. A local/remote flag would misprice a free self-hosted server.

# s1a/decision_models/base.py
bills_input_tokens: bool  # True when input tokens are priced at JEV_USD_PER_INPUT_TOKEN

# s1a/decision_models/jev.py
bills_input_tokens = True

# laya.py (LayaModel), cua.py (CuaS1Model), baselines.py (RandomModel, RuleModel), fakes.py (ScriptedModel)
bills_input_tokens = False

The tool and browser fronts have the same string check. Switch all three:

  • s1a/rails.py:124: ... if decision_model.bills_input_tokens else 0
  • s1a/browser/decision_model.py:778: ... if self._decision_model.bills_input_tokens else 0
  • s1a/tool/loop.py:277: ToolDecisionModel wraps the decision model. Copy the flag in its __init__ next to self.name (s1a/tool/models.py:82), then gate the sum on isinstance(model, ToolDecisionModel) and model.bills_input_tokens. Keep excluding the "llm" fallback ticks (loop.py:208); their tokens are already priced as chat tokens.

Python doesn't enforce a bare annotation. A backend that forgets the flag raises AttributeError at its first summary. A smoke test catches that at test time:

# tests/test_decision_models_base.py
def test_every_decision_model_declares_whether_it_bills_input_tokens(self) -> None:
    backends = {backend.__name__: backend for backend in DecisionModel.__subclasses__()}
    self.assertLessEqual({"JevModel", "LayaModel", "CuaS1Model", "RandomModel", "RuleModel", "ScriptedModel"}, set(backends))
    for name, backend in backends.items():
        with self.subTest(backend=name):
            self.assertIsInstance(vars(backend).get("bills_input_tokens"), bool)

__subclasses__() only sees imported modules. The subset check fails if an import is missing, so import the backend modules at the top of the test file.

With the flag in place, the comment at rails.py:123 can go. Your two new tests should pass unchanged and cover the refactor.

@alexwong10
alexwong10 force-pushed the fix/rail-local-model-costs branch from 091b097 to 68f29f4 Compare September 28, 2026 10:36
@alexwong10 alexwong10 changed the title [Fix] Avoid charging Jev API prices for local rail models [Fix] Let decision backends declare billable input tokens Sep 28, 2026
@alexwong10

Copy link
Copy Markdown
Contributor Author

Hi maintainers, all CI jobs that ran on this PR have passed. The merge box still lists core (3.11) and core (3.13) as expected required checks, while the workflow now reports core (ubuntu-latest, 3.11) and core (ubuntu-latest, 3.13).

Could someone with repository settings access update the required check names? An approving review is also still required once the changes look good. Thanks!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants