[Fix] Let decision backends declare billable input tokens - #17
alexwong10 wants to merge 2 commits into
Conversation
There was a problem hiding this comment.
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 = FalseThe tool and browser fronts have the same string check. Switch all three:
s1a/rails.py:124:... if decision_model.bills_input_tokens else 0s1a/browser/decision_model.py:778:... if self._decision_model.bills_input_tokens else 0s1a/tool/loop.py:277:ToolDecisionModelwraps the decision model. Copy the flag in its__init__next toself.name(s1a/tool/models.py:82), then gate the sum onisinstance(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.
091b097 to
68f29f4
Compare
|
Hi maintainers, all CI jobs that ran on this PR have passed. The merge box still lists Could someone with repository settings access update the required check names? An approving review is also still required once the changes look good. Thanks! |
Why
A rail run with
--model layareported 300 local input tokens as $0.000013 of Jev API charges. Checkingmodel.name == "jev"to decide billing would also miss a future paid backend with a different name.How
DecisionModel.bills_input_tokensis a required backend declaration. Jev opts in; Laya, Cua, Random, Rule, and Scripted opt out. Rail, browser, and tool reports read the flag, andToolDecisionModelcarries it through its wrapper. Tool accounting excludesllmfallback ticks, which are already priced as chat tokens. A true flag uses the existingJEV_USD_PER_INPUT_TOKENrate.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
scriptedin each front, and an explicit billing declaration on every built-in backend.Verification
pytest -q— 495 passed, 41 skipped, 80 subtests passed.ruff format --check .,ruff check .,ty check,uv lock --check --offline,uv build --offline, andscripts/smoke.shin Git Bash.