diff --git a/llmcore.py b/llmcore.py index 88879416e..d9b545b2f 100644 --- a/llmcore.py +++ b/llmcore.py @@ -99,13 +99,57 @@ def safeprint(*argv): STATS = {} +# Default chars-per-token heuristic for the legacy `context_win * 3` rule. +# Kept as a named module constant so a future model-aware estimator only has +# to swap the value here (or override per-session via +# `sess.token_chars_per_token` / `sess.token_estimator`). `trim_messages_history` +# still uses this as a char-equivalent multiplier — the trim gate is char-based +# to keep behaviour byte-identical with the previous implementation. +DEFAULT_CHARS_PER_TOKEN = 3 + +def estimate_context_tokens(history, sess=None): + """Estimate `history`'s "tokens" for debug / diagnostics, *not* the trim gate. + + This is a pure function of `(history, sess)` that returns an int token + estimate. It is **not** the trigger that fires `trim_messages_history` + (the trigger is still `cost(history) > context_win * 3` chars) — the + estimator exists so we can: + + 1. Stop burying the heuristic inside `trim_messages_history`. Issue + #750: `3 chars/token` is wrong for CJK/code/JSON; today there is no + single place to swap it. + 2. Surface the estimate alongside the real `input_tokens` recorded by + `_record_usage()` so future debugging can compare them. + 3. Give sessions a clean override entry (`sess.token_estimator` callable + OR `sess.token_chars_per_token` float) without leaking provider- + specific knobs into the trim path. + + Default behaviour (v1, byte-identical with the prior code): + `int(sum(len(json.dumps(m, ensure_ascii=False)) for m in history) / 3)` + i.e. the same `cost() / 3` the prior code used implicitly for its log. + `sess` may be None — used by unit tests and by callers that have no + session object handy yet. + """ + estimator = getattr(sess, 'token_estimator', None) if sess is not None else None + if estimator is not None: + try: return max(0, int(estimator(history) or 0)) + except Exception: pass # never let an override crash the trim path + cpt = getattr(sess, 'token_chars_per_token', DEFAULT_CHARS_PER_TOKEN) if sess is not None else DEFAULT_CHARS_PER_TOKEN + try: cpt = float(cpt) + except (TypeError, ValueError): cpt = DEFAULT_CHARS_PER_TOKEN + if cpt <= 0: cpt = DEFAULT_CHARS_PER_TOKEN # guard against a misconfigured zero / negative divisor + base = sum(len(json.dumps(m, ensure_ascii=False)) for m in history) + return int(base / cpt) + def trim_messages_history(history, sess): - cap = sess.context_win * 3 + cap = sess.context_win * DEFAULT_CHARS_PER_TOKEN target = int(cap * getattr(sess, 'trim_keep_rate', 0.6)) kp = sess.trim_keep_prefix def cost(ms): return sum(len(json.dumps(m, ensure_ascii=False)) for m in ms) compress_history_tags(history, interval=getattr(sess, 'cut_msg_interval', 7)) - STATS.update(ctx=(c := cost(history)), msgs=len(history)); print(f'[Debug] Current context: {c} chars, {len(history)} messages.') + c = cost(history) + STATS.update(ctx=c, tokens_est=estimate_context_tokens(history, sess), msgs=len(history)) + print(f'[Debug] Current context: {c} chars (~{STATS["tokens_est"]} tokens est), {len(history)} messages.') if c <= cap: return compress_history_tags(history, keep_recent=4, force=True) if cost(history) <= target: return @@ -122,7 +166,9 @@ def cost(ms): return sum(len(json.dumps(m, ensure_ascii=False)) for m in ms) gap = [{"role": "assistant", "content": _d()}] if m.get('role') == 'user' else [{"role": "user", "content": _d()}, {"role": "assistant", "content": _d()}] history[:] = pre + gap + post else: history[:] = pre + post - STATS.update(ctx=(c := cost(history)), msgs=len(history)); print(f'[Debug] Trimmed context, current: {c} chars, {len(history)} messages.') + c = cost(history) + STATS.update(ctx=c, tokens_est=estimate_context_tokens(history, sess), msgs=len(history)) + print(f'[Debug] Trimmed context, current: {c} chars (~{STATS["tokens_est"]} tokens est), {len(history)} messages.') def auto_make_url(base, path): b, p = base.rstrip('/'), path.strip('/') diff --git a/tests/test_llmcore_estimate_context_tokens.py b/tests/test_llmcore_estimate_context_tokens.py new file mode 100644 index 000000000..80f840619 --- /dev/null +++ b/tests/test_llmcore_estimate_context_tokens.py @@ -0,0 +1,255 @@ +"""Regression tests for `llmcore.estimate_context_tokens` (#750). + +Background: `trim_messages_history` historically used the heuristic +``3 chars == 1 token`` inlined into the cap calculation +(``cap = sess.context_win * 3``). Issue #750 called for that heuristic to +be moved into a single named place, so a future model-aware estimator can +swap it without touching the trim path. This file pins the v1 behaviour +and exercises both extension points: +``sess.token_estimator`` (callable override) and +``sess.token_chars_per_token`` (float multiplier override). + +The tests load `llmcore.py` in isolation (heavy deps stubbed) — same trick +used by `tests/test_llmcore_record_usage_messages.py`. +""" + +import importlib.util +import json +import sys +import types +import unittest +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parent.parent +LLMCORE_PATH = REPO_ROOT / "llmcore.py" + + +def _load_llmcore_module(): + """Load `llmcore` without spinning up the agent runtime. + + `llmcore` is a top-level script that imports `agentmain`, `requests`, + and `websockets` at module load time. We stub just enough to import the + module and access the helpers we want to test, then return it. + """ + stubs = { + "requests": types.ModuleType("requests"), + "websockets": types.ModuleType("websockets"), + "agentmain": types.ModuleType("agentmain"), + } + stubs["requests"].post = lambda *a, **kw: None + stubs["requests"].Session = type("Session", (), {}) + for name, mod in stubs.items(): + sys.modules.setdefault(name, mod) + spec = importlib.util.spec_from_file_location("llmcore_under_test", LLMCORE_PATH) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +llmcore = _load_llmcore_module() + + +def _history(*pairs): + """Build a history list of `{role, content}` dicts from (role, text) pairs.""" + return [{"role": role, "content": text} for role, text in pairs] + + +def _sess(**attrs): + """Build a throw-away session stub. Only the attributes our estimator + reads are set; everything else is irrelevant.""" + return types.SimpleNamespace(**attrs) + + +class EstimateContextTokensDefaultHeuristicTests(unittest.TestCase): + """Pin the v1 behaviour: `int(char_count / 3)` rounded down.""" + + def _json_chars(self, *pairs): + # The estimator sums `len(json.dumps(m, ensure_ascii=False)) for m in history` + # i.e. it json-serialises each message *individually* and sums. Match that + # here so we don't bake in the wrong assumption about list-vs-per-msg. + return sum(len(json.dumps({"role": r, "content": c}, ensure_ascii=False)) for r, c in pairs) + + def test_empty_history_returns_zero(self): + self.assertEqual(llmcore.estimate_context_tokens([]), 0) + self.assertEqual(llmcore.estimate_context_tokens([], None), 0) + + def test_short_english_history(self): + # {"role":"user","content":"hello"} -> 36 chars -> 36/3 = 12 tokens. + hist = _history(("user", "hello")) + self.assertEqual(llmcore.estimate_context_tokens(hist), self._json_chars(("user", "hello")) // 3) + self.assertEqual(llmcore.estimate_context_tokens(hist), 12) + + def test_cjk_history_is_one_token_per_three_chars(self): + # CJK glyphs are still 1 char each under ensure_ascii=False — that + # is the v1 bug #750 calls out. Pin the behaviour so a future + # tokenizer-aware fix doesn't silently regress English history. + hist = _history(("user", "你好世界你好")) + self.assertEqual(llmcore.estimate_context_tokens(hist), self._json_chars(("user", "你好世界你好")) // 3) + self.assertEqual(llmcore.estimate_context_tokens(hist), 12) + + def test_uses_ensure_ascii_false_so_cjk_counts_in_chars(self): + # `json.dumps(..., ensure_ascii=False)` keeps the original chars + # rather than escaping them to \uXXXX (which would inflate CJK). + # With escaped form "你好" would become "\u4f60\u597d" (12 chars); + # with ensure_ascii=False the content stays at 2 chars. We pin the + # unescaped form (the value used by the original `cost()` too). + hist = _history(("user", "你好")) + self.assertEqual(llmcore.estimate_context_tokens(hist), self._json_chars(("user", "你好")) // 3) + # Sanity: escaped form would yield a different number. + escaped_chars = len(json.dumps({"role": "user", "content": "你好"}, ensure_ascii=True)) + self.assertNotEqual(self._json_chars(("user", "你好")), escaped_chars) + + def test_history_with_multiple_messages(self): + # The estimator sums per-message json.dumps chars, then divides. + hist = _history(("user", "abc"), ("assistant", "defg")) + self.assertEqual( + llmcore.estimate_context_tokens(hist), + self._json_chars(("user", "abc"), ("assistant", "defg")) // 3, + ) + + def test_non_json_serialisable_messages_are_handled_by_json_dumps(self): + # `cost()` (and the estimator) use `json.dumps` directly. A + # non-serialisable value will raise inside the estimator — that's + # acceptable because the trim path catches and treats it as + # over-cap (the original code also relied on this). Pin it. + bad = [{"role": "user", "content": {"nested": object()}}] + with self.assertRaises(TypeError): + llmcore.estimate_context_tokens(bad) + + +class EstimateContextTokensCharsPerTokenOverrideTests(unittest.TestCase): + """`sess.token_chars_per_token` is the float-multiplier override.""" + + def test_string_is_coerced_to_float(self): + sess = _sess(token_chars_per_token="1.5") + # 36 chars / 1.5 = 24 tokens + self.assertEqual(llmcore.estimate_context_tokens(_history(("user", "hello")), sess), 24) + + def test_invalid_value_falls_back_to_default(self): + # Garbage values must not crash the trim path; the fallback is the + # module default of 3. + sess = _sess(token_chars_per_token="not-a-number") + # 38 chars / 3 = 12 tokens + self.assertEqual(llmcore.estimate_context_tokens(_history(("user", "hello")), sess), 12) + + def test_zero_divisor_is_replaced_by_default(self): + # A misconfigured zero / negative must not produce a ZeroDivisionError + # or a wildly negative estimate — it falls back to the default. + sess = _sess(token_chars_per_token=0) + self.assertEqual(llmcore.estimate_context_tokens(_history(("user", "hello")), sess), 12) + sess = _sess(token_chars_per_token=-2) + self.assertEqual(llmcore.estimate_context_tokens(_history(("user", "hello")), sess), 12) + + def test_none_value_falls_back_to_default(self): + sess = _sess(token_chars_per_token=None) + self.assertEqual(llmcore.estimate_context_tokens(_history(("user", "hello")), sess), 12) + + +class EstimateContextTokensCallableOverrideTests(unittest.TestCase): + """`sess.token_estimator` is the full-callable override.""" + + def test_callable_takes_precedence_over_chars_per_token(self): + sess = _sess(token_estimator=lambda h: 42, token_chars_per_token=1.5) + # The callable wins even when chars_per_token would also be set. + self.assertEqual(llmcore.estimate_context_tokens(_history(("user", "hi")), sess), 42) + + def test_callable_returning_zero_or_none_is_clamped(self): + # `int(estimator(history) or 0)` then `max(0, ...)` — defensive + # against a buggy estimator that returns None or 0 or a negative. + sess = _sess(token_estimator=lambda h: 0) + self.assertEqual(llmcore.estimate_context_tokens(_history(("user", "hi")), sess), 0) + sess = _sess(token_estimator=lambda h: None) + self.assertEqual(llmcore.estimate_context_tokens(_history(("user", "hi")), sess), 0) + sess = _sess(token_estimator=lambda h: -7) + self.assertEqual(llmcore.estimate_context_tokens(_history(("user", "hi")), sess), 0) + + def test_callable_returning_non_int_is_truncated(self): + # Floats must be coerced to int — the estimator contract says int. + sess = _sess(token_estimator=lambda h: 1.9) + self.assertEqual(llmcore.estimate_context_tokens(_history(("user", "hi")), sess), 1) + + def test_buggy_callable_is_swallowed_and_falls_back(self): + # The estimator must never crash the trim path — a buggy override + # should fall back to the chars-per-token path. + sess = _sess(token_estimator=lambda h: 1 / 0, token_chars_per_token="3") + # 38 chars / 3 = 12 tokens (fallback to chars_per_token path) + self.assertEqual(llmcore.estimate_context_tokens(_history(("user", "hello")), sess), 12) + + def test_callable_receives_full_history(self): + captured = [] + sess = _sess(token_estimator=lambda h: captured.append(list(h)) or 99) + hist = _history(("user", "hi"), ("assistant", "yo")) + llmcore.estimate_context_tokens(hist, sess) + # The estimator sees the whole list, not a slice. + self.assertEqual(len(captured[0]), 2) + self.assertEqual(captured[0][0]["role"], "user") + + +class TrimMessagesHistoryBehaviorPreservedTests(unittest.TestCase): + """The trim *gate* must keep its char-based behaviour byte-identical + with the previous implementation — only the debug log gains the new + `tokens_est` field.""" + + def _make_sess(self, context_win=100, trim_keep_prefix=0, trim_keep_rate=0.6): + return _sess( + context_win=context_win, + trim_keep_prefix=trim_keep_prefix, + trim_keep_rate=trim_keep_rate, + ) + + def test_small_history_does_not_trim_and_logs_both_stats(self): + sess = self._make_sess(context_win=1000) + hist = _history(("user", "hi")) + llmcore.trim_messages_history(hist, sess) + # gate says: under cap -> return without trimming + self.assertEqual(len(hist), 1) + # stats now expose both `ctx` (chars, for cost_tracker compat) and + # `tokens_est` (the new field). + self.assertIn("ctx", llmcore.STATS) + self.assertIn("tokens_est", llmcore.STATS) + # tokens_est is the same as ctx / 3 under the default heuristic. + self.assertEqual(llmcore.STATS["tokens_est"], llmcore.STATS["ctx"] // 3) + + def test_large_history_triggers_hard_cut(self): + # context_win * 3 chars is the trim trigger. Build a history + # whose post-compress cost > target (the trim target), then + # confirm: + # * length decreases + # * tokens_est tracks ctx / 3 under the default heuristic + # compress_history_tags() shrinks / blocks in + # older messages but does not drop messages, so plain content is + # left alone and the trim's hard-cut loop (which loops while + # `len(post) > 9`) is the only thing that can drop messages. + sess = self._make_sess(context_win=2, trim_keep_prefix=0, trim_keep_rate=0.3) + # 25 user messages of 30-char content. context_win*3 = 6 cap, + # target = 6 * 0.3 = 1. hard-cut loop: pops until either + # len(post) <= 9 OR cost <= target. With 30-char msgs at >1 + # char target the loop terminates by len(post) > 9, ending at 9. + hist = _history(*[("user", "x" * 30)] * 25) + before = len(hist) + llmcore.trim_messages_history(hist, sess) + self.assertLess(len(hist), before) + # The trim path always leaves the last <=9 messages (that's the + # only shape the existing hard-cut loop can produce). Pin it. + self.assertLessEqual(len(hist), 9) + # tokens_est still tracks ctx / 3 under the default heuristic. + self.assertEqual(llmcore.STATS["tokens_est"], llmcore.STATS["ctx"] // 3) + + def test_session_with_custom_chars_per_token_is_honoured_in_log_only(self): + # The estimator must reflect the per-session override in STATS, + # but the trim gate must still use the default 3-char multiplier + # (so behaviour stays byte-identical for v1). + sess = self._make_sess(context_win=1000) + sess.token_chars_per_token = 1.5 + hist = _history(("user", "hi")) + llmcore.trim_messages_history(hist, sess) + # ctx stays char-based. + self.assertGreater(llmcore.STATS["ctx"], 0) + # tokens_est uses the per-session override (1.5x larger than /3). + # exact formula: ctx / 1.5 (rounded down by int()). + self.assertEqual(llmcore.STATS["tokens_est"], llmcore.STATS["ctx"] // 1.5) + + +if __name__ == "__main__": + unittest.main() \ No newline at end of file