claude-code
d778c113e9
werkator build succeeded after 01:14
werkrevisor 0.13.0: mechanical file changes skip the model. - A file is mechanical only if every changed line is a pure rename, a whitespace change (not in whitespace-significant files) or an import line of a known language; it goes neither into the diff nor into the full files for the model. - A PR of only mechanical files is GREEN-eligible without a model call; hard rules still apply. - Additional changes: the example host configs match mih09 (`mode = "combined"`, `token_count_factor = 1.55`); an unknown change kind in a changed `.werkrevisor.toml` gives YELLOW, not RED. Details: the PR document under `docs/prs/`. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Reviewed-on: #11 Co-authored-by: claude-code <claude-code@javagil.de>
162 lines
9.2 KiB
Python
162 lines
9.2 KiB
Python
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()
|