import json import tomllib import unittest from unittest import mock from tests.fakes import ROOT, FakeGitea, FakeKey, FakeModel, config, deps, diff_for, gitea_file, pr_ref from werkrevisor import budget from werkrevisor.cli import execute from werkrevisor.config import ConfigError, build_config from werkrevisor.redact import Redactor from werkrevisor.verdict import Verdict DOC = "doc/PR/1.md" def words(text): """A fake tokenizer: one token per padding word of FakeModel, nothing else.""" return text.count("word ") def run(model, cfg=None, key=None): gitea = FakeGitea(files=[gitea_file(DOC)], diff=diff_for(DOC, added=["Some text."])) d = deps(gitea, model, key) outcome = execute(cfg or config(gate={"yellow_gate": "block"}), pr_ref(), d, Redactor()) return outcome, gitea, json.loads(d.out.getvalue().strip()) class TokenCountTest(unittest.TestCase): def test_count_uses_the_tokenizer_and_the_factor(self): for factor, expected in ((1.15, 1150), (1, 1000), (2, 2000)): with self.subTest(factor=factor): cfg = config(model={"token_count_factor": factor}) with mock.patch.object(budget, "token_counter", return_value=words): count = budget.count_prompt("word " * 400, "word " * 600, cfg.model, lambda: self.fail("no estimate with a tokenizer")) self.assertEqual((count.raw_tokens, count.input_tokens, count.method, count.factor), (1000, expected, "o200k_base", factor)) self.assertEqual(count.worst_case_usd, budget.worst_case_usd(expected, cfg.model)) def test_real_o200k_base_tokenizer_loads_offline(self): import tiktoken counter = budget.token_counter() self.assertIsNotNone(counter, "litellm's bundled o200k_base must load without network") text = "def f(n):\n return 1 / n # <|endoftext|> stays plain text\n" self.assertEqual(counter(text), len(tiktoken.get_encoding("o200k_base").encode(text, disallowed_special=()))) count = budget.count_prompt("system", text, config().model, lambda: self.fail("no estimate")) self.assertEqual(count.raw_tokens, counter("system") + counter(text)) def test_fallback_to_the_estimate_is_flagged(self): cfg = config() est = budget.estimate(3500, 0, cfg.model, cfg.budget) with mock.patch.object(budget, "token_counter", return_value=None): count = budget.count_prompt("system", "user", cfg.model, lambda: est) self.assertEqual((count.input_tokens, count.worst_case_usd, count.method, count.raw_tokens), (est.input_tokens, est.worst_case_usd, "estimated", None)) self.assertIn("token count estimated", count.describe()) class TokenCountFlowTest(unittest.TestCase): def test_limit_is_decided_on_the_count_times_the_factor(self): cfg = config(model={"max_context_tokens": 1000, "token_count_factor": 1.15}, gate={"yellow_gate": "block"}) # 869 × 1.15 = 999.35 fits into 1000 tokens, 870 × 1.15 = 1000.5 does not. for padding, called in ((869, True), (870, False)): with self.subTest(padding=padding): model, key = FakeModel(prompt_files=[DOC], prompt_words=padding), FakeKey() with mock.patch.object(budget, "token_counter", return_value=words): outcome, _, record = run(model, cfg, key) self.assertEqual((bool(model.calls), key.calls), (called, 2 if called else 0)) self.assertEqual(record["counted_raw_tokens"], padding) codes = [r.code for r in outcome.result.reasons] self.assertEqual("too_large_tokens" in codes, not called) def test_estimated_count_is_named_in_reason_comment_and_usage(self): cfg = config(model={"max_context_tokens": 1000}, gate={"yellow_gate": "block"}) model = FakeModel(prompt_files=[DOC]) with mock.patch.object(budget, "token_counter", return_value=None): outcome, gitea, record = run(model, cfg) self.assertIs(outcome.result.verdict, Verdict.YELLOW) self.assertEqual(model.calls, []) reason = next(r for r in outcome.result.reasons if r.code == "too_large_tokens") self.assertIn("(token count estimated: the tokenizer could not load)", reason.text) self.assertIn("token count estimated: about", gitea.created_comments[0]) self.assertEqual((record["token_count_method"], record["counted_raw_tokens"]), ("estimated", None)) def test_usage_record_has_count_and_reported_prompt_tokens(self): model = FakeModel(prompt_files=[DOC]) outcome, gitea, record = run(model) system, user = model.prompts[0] counter = budget.token_counter() raw = counter(system) + counter(user) self.assertEqual(record["counted_raw_tokens"], raw) self.assertEqual(record["counted_input_tokens"], outcome.result.token_count.input_tokens) self.assertEqual((record["token_count_factor"], record["token_count_method"]), (1.55, "o200k_base")) self.assertEqual(record["worst_case_usd"], outcome.result.token_count.worst_case_usd) self.assertEqual(record["prompt_tokens"], 1000) self.assertIn(f"tokens counted before the call {record['counted_input_tokens']}", gitea.created_comments[0]) def test_no_prompt_no_count(self): # Assembled, so that werkrevisor's own secret scan does not skip the model review of this file. private_key_marker = "-----BEGIN " + "PRIVATE KEY-----" gitea = FakeGitea(files=[gitea_file(DOC)], diff=diff_for(DOC, added=[private_key_marker])) d = deps(gitea, FakeModel(prompt_files=[DOC])) execute(config(gate={"yellow_gate": "block"}), pr_ref(), d, Redactor()) record = json.loads(d.out.getvalue().strip()) self.assertEqual((record["counted_input_tokens"], record["token_count_method"]), (None, None)) class RetryTest(unittest.TestCase): def test_retry_with_the_same_prompt_queries_the_key_only_once(self): # The second key query would fail; it must not happen, because the first call already cost money. model, key = FakeModel(prompt_files=[DOC], prompt_words=[1000, 1000]), FakeKey(remaining=7.0, fail_from_call=2) outcome, _, record = run(model, key=key) self.assertIs(outcome.result.verdict, Verdict.GREEN) self.assertEqual((len(model.prompts), key.calls), (2, 2)) self.assertEqual(record["budget_remaining_before_usd"], 7.0) self.assertEqual(outcome.result.llm_status, "done") def test_retry_with_a_grown_prompt_above_the_limit_stops_and_keeps_the_usage(self): cfg = config(model={"max_context_tokens": 1000, "token_count_factor": 1.15}, gate={"yellow_gate": "block"}) model, key = FakeModel(prompt_files=[DOC], prompt_words=[800, 900]), FakeKey() with mock.patch.object(budget, "token_counter", return_value=words): outcome, gitea, record = run(model, cfg, key) self.assertEqual((model.calls, key.calls), ([], 1)) self.assertEqual(record["counted_raw_tokens"], 900) self.assertIn("too_large_tokens", [r.code for r in outcome.result.reasons]) self.assertEqual((record["calls"], record["prompt_tokens"], record["cost_status"]), (1, 1000, "estimated")) self.assertGreater(record["cost_usd"], 0) self.assertNotIn("no model call, cost 0 USD", gitea.created_comments[0]) self.assertEqual(outcome.result.llm_status, "stopped after 1 model call(s): skipped (too large)") def test_retry_with_a_smaller_prompt_keeps_the_first_decision(self): model, key = FakeModel(prompt_files=[DOC], prompt_words=[900, 800]), FakeKey() with mock.patch.object(budget, "token_counter", return_value=words): outcome, _, record = run(model, key=key) self.assertEqual((len(model.calls), key.calls, record["counted_raw_tokens"]), (1, 2, 900)) def test_stop_without_an_earlier_call_costs_nothing(self): cfg = config(model={"max_context_tokens": 1000}, gate={"yellow_gate": "block"}) with mock.patch.object(budget, "token_counter", return_value=words): outcome, _, record = run(FakeModel(prompt_files=[DOC], prompt_words=[900]), cfg) self.assertEqual((record["cost_usd"], record["cost_status"], record["calls"]), (0.0, "none", 0)) class TokenCountConfigTest(unittest.TestCase): def test_factor_must_be_positive(self): for value in (0, -1, True, "1.15", float("inf"), float("nan")): with self.subTest(value=value): with self.assertRaisesRegex(ConfigError, r"\[model\] token_count_factor must be a positive number"): config(model={"token_count_factor": value}) self.assertEqual(config().model.token_count_factor, 1.55) self.assertEqual(config(model={"token_count_factor": 1}).model.token_count_factor, 1) def test_factor_is_host_only(self): with open(f"{ROOT}/config/werkprobe.toml", "rb") as f: host = tomllib.load(f) with self.assertRaisesRegex(ConfigError, r"must not set \[model\] token_count_factor: only the host config"): build_config(host, {"model": {"token_count_factor": 0.5}}) if __name__ == "__main__": unittest.main()