diff --git a/strix/llm/context_budget.py b/strix/llm/context_budget.py index 47640184..6bd805e3 100644 --- a/strix/llm/context_budget.py +++ b/strix/llm/context_budget.py @@ -68,10 +68,18 @@ def output_limit(model: str) -> int: def count_tokens(model: str, text: str) -> int: - """Token count for ``text`` under ``model`` (chars/4 fallback).""" + """Token count for ``text`` under ``model``. + + LiteLLM's counter handles known tokenizers (and defaults to a tiktoken + encoding otherwise). If it still can't count, fall back to a *conservative* + estimate: token density varies, and dense text (code, base64, CJK) can run + well under 4 chars/token, so we assume ~3 to over-estimate rather than + under-estimate — an under-estimate would let a summary request be packed + past the real context window and get rejected. + """ if not text: return 0 try: return int(litellm.token_counter(model=_lookup_key(model), text=text)) except Exception: # noqa: BLE001 - tokenizer may be unavailable for some models. - return len(text) // 4 + return -(-len(text) // 3) diff --git a/tests/test_context_budget.py b/tests/test_context_budget.py index 420d35b4..b2032137 100644 --- a/tests/test_context_budget.py +++ b/tests/test_context_budget.py @@ -38,8 +38,10 @@ def test_count_tokens_fallback_on_error(monkeypatch: pytest.MonkeyPatch) -> None raise RuntimeError("no tokenizer") monkeypatch.setattr("strix.llm.context_budget.litellm.token_counter", _raise) - text = "x" * 400 - assert context_budget.count_tokens("weird-model", text) == 100 + # Conservative ~3-chars/token estimate, rounded up, so budgets never + # under-count when no tokenizer is available. + assert context_budget.count_tokens("weird-model", "x" * 400) == 134 + assert context_budget.count_tokens("weird-model", "x") == 1 def test_count_tokens_empty_is_zero() -> None: