diff --git a/src/granite_switch_decision/plugin.py b/src/granite_switch_decision/plugin.py index b85c325..2b56b48 100644 --- a/src/granite_switch_decision/plugin.py +++ b/src/granite_switch_decision/plugin.py @@ -443,18 +443,20 @@ def _enforce_context_limits( ) state_text = _text_of(state) if state not in ("", None, {}, []) else "(empty)" state_tokens = _count_tokens(tok, state_text) - total = 0 + # jev: a request is the state plus all its questions; a question is the state plus that question. + total = state_tokens for name, q in questions.items(): options = q.criteria.items() if isinstance(q.criteria, dict) else enumerate(q.criteria or []) option_text = "\n".join(f"{key} {_text_of(desc)}" for key, desc in options) - n = state_tokens + _count_tokens(tok, f"{_text_of(q.instructions)}\n{option_text}") + question_tokens = _count_tokens(tok, f"{_text_of(q.instructions)}\n{option_text}") + n = state_tokens + question_tokens if n > per_question_cap: raise ValueError( f"question {name!r} is approximately {n} tokens, exceeding the " f"{per_question_cap}-token maximum context length per question " f"(max_tokens_exceeded)" ) - total += n + total += question_tokens if total > total_cap: raise ValueError( f"request is approximately {total} tokens total, exceeding the " diff --git a/tests/test_plugin.py b/tests/test_plugin.py index 7377adf..550b8b5 100644 --- a/tests/test_plugin.py +++ b/tests/test_plugin.py @@ -430,6 +430,22 @@ def test_enforce_context_limits_total_message_matches_a_capacity_marker(): assert any(marker in str(exc_info.value) for marker in _CAPACITY_MARKERS) +def test_enforce_context_limits_counts_the_state_once_per_request(): + """jev counts a request as the state plus all questions, not the state once per question. + 37 short questions over a 5,000-token state are about 5,200 tokens, well under the 64k cap.""" + tok = _FakeTok() + questions = {f"q{i}": sp.Question(type="noul", instructions="is it urgent now ok") for i in range(37)} + sp._enforce_context_limits(tok, 131072, False, "state " * 5000, questions) # must not raise + + +def test_enforce_context_limits_still_rejects_a_request_that_is_over_the_cap(): + """Each question fits (20,000 + 10,000), but the state plus all five is 70,000 tokens.""" + tok = _FakeTok() + questions = {f"q{i}": sp.Question(type="noul", instructions="word " * 10_000) for i in range(5)} + with pytest.raises(ValueError, match="per request"): + sp._enforce_context_limits(tok, 131072, False, "state " * 20_000, questions) + + def test_enforce_context_limits_accepts_small_request(): tok = _FakeTok() q = sp.Question(type="noul", instructions="a short question")