From 92736639b10b8fc07eab4e44cc8eb00c9c617ba5 Mon Sep 17 00:00:00 2001 From: "detail-app[bot]" <180357370+detail-app[bot]@users.noreply.github.com> Date: Sun, 6 Sep 2026 10:39:41 +0000 Subject: [PATCH 1/2] fix(dspy): reconcile max_tokens override with aliased max_completion_tokens for reasoning models --- dspy/clients/lm.py | 16 +++++++- tests/clients/test_lm.py | 83 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 98 insertions(+), 1 deletion(-) diff --git a/dspy/clients/lm.py b/dspy/clients/lm.py index b6fa8ffb4a..9ffdfa773f 100644 --- a/dspy/clients/lm.py +++ b/dspy/clients/lm.py @@ -401,6 +401,18 @@ def infer_provider(self) -> Provider: return OpenAIProvider() return Provider() + def copy(self, **kwargs): + """Return a copy of the language model with updated parameters. + + Reconciles the user-facing ``max_tokens`` override with the + ``max_completion_tokens`` alias stored for OpenAI reasoning models, so + the copied LM matches what ``LM.__init__`` would have produced for the + same ``max_tokens`` value and the override survives ``dump_state()``. + """ + if "max_tokens" in kwargs and _is_openai_reasoning_model(self.model): + kwargs["max_completion_tokens"] = kwargs.pop("max_tokens") + return super().copy(**kwargs) + def dump_state(self): """Return a sanitized reconstruction state for this LM. @@ -419,7 +431,9 @@ def dump_state(self): if self.use_developer_role: state["use_developer_role"] = self.use_developer_role if _is_openai_reasoning_model(self.model) and "max_completion_tokens" in state: - state["max_tokens"] = state.pop("max_completion_tokens") + if "max_tokens" not in state: + state["max_tokens"] = state["max_completion_tokens"] + state.pop("max_completion_tokens") return state @classmethod diff --git a/tests/clients/test_lm.py b/tests/clients/test_lm.py index 6588348092..169ce03465 100644 --- a/tests/clients/test_lm.py +++ b/tests/clients/test_lm.py @@ -887,6 +887,89 @@ class CustomLM(dspy.BaseLM): assert lm.kwargs == {"temperature": 0.1, "max_tokens": None} +def test_copy_max_tokens_reasoning_model_preserves_budget(): + lm = dspy.LM("openai/gpt-5", temperature=1.0, max_tokens=16000, api_key="sk-fake") + lm2 = lm.copy(max_tokens=32000) + + # copy() should update the aliased key, not insert a second one + token_keys = [k for k in lm2.kwargs if "token" in k] + assert token_keys == ["max_completion_tokens"], f"unexpected token keys: {token_keys}" + assert lm2.kwargs["max_completion_tokens"] == 32000 + + # dump_state() must preserve the user's raised budget + assert lm2.dump_state()["max_tokens"] == 32000 + + +def test_copy_max_tokens_reasoning_model_chat_request_has_single_token_budget(): + """On the chat path, copy(max_tokens=...) must send only max_completion_tokens to litellm — no stale dual key.""" + lm = dspy.LM("openai/gpt-5", temperature=1.0, max_tokens=16000, cache=False) + copied = lm.copy(max_tokens=32000) + + with mock.patch("dspy.clients.lm.litellm_completion", return_value=_model_response("hi")) as completion: + copied("query") + + request = completion.call_args.kwargs["request"] + assert request["max_completion_tokens"] == 32000 + assert "max_tokens" not in request + + +def test_copy_max_tokens_reasoning_model_responses_request_has_single_token_budget(): + """On the responses path, copy(max_tokens=...) must not leak stale max_completion_tokens next to max_output_tokens.""" + lm = dspy.LM("openai/gpt-5", model_type="responses", temperature=1.0, max_tokens=16000, cache=False) + copied = lm.copy(max_tokens=32000) + + with mock.patch("litellm.responses", return_value=make_response([])) as responses: + copied("query") + + call_kwargs = responses.call_args.kwargs + assert call_kwargs["max_output_tokens"] == 32000 + assert "max_completion_tokens" not in call_kwargs + assert "max_tokens" not in call_kwargs + + +def test_copy_max_tokens_reasoning_model_round_trips_via_dump_load_state(): + lm = dspy.LM("openai/gpt-5", temperature=1.0, max_tokens=16000, cache=False) + copied = lm.copy(max_tokens=32000) + + state = copied.dump_state() + assert state["max_tokens"] == 32000 + assert "max_completion_tokens" not in state + + reloaded = dspy.LM.load_state(state) + assert reloaded.kwargs["max_completion_tokens"] == 32000 + assert reloaded.dump_state()["max_tokens"] == 32000 + + +def test_copy_max_tokens_none_reasoning_model_drops_token_budget(): + """copy(max_tokens=None) follows the standard copy() "None means remove" semantics for the aliased key.""" + lm = dspy.LM("openai/gpt-5", temperature=1.0, max_tokens=16000, cache=False) + copied = lm.copy(max_tokens=None) + + assert "max_completion_tokens" not in copied.kwargs + assert "max_tokens" not in copied.kwargs + + +def test_copy_max_tokens_non_reasoning_model_preserves_max_tokens_key(): + """Non-reasoning models keep using max_tokens directly (no reasoning-model aliasing).""" + lm = dspy.LM("openai/gpt-4o", max_tokens=100, cache=False) + copied = lm.copy(max_tokens=200) + + assert copied.kwargs["max_tokens"] == 200 + assert "max_completion_tokens" not in copied.kwargs + + +def test_dump_state_preserves_existing_max_tokens_when_dual_keys_present(): + """Defense in depth: if a dual-key state ever exists, dump_state keeps the existing max_tokens and drops the stale alias.""" + lm = dspy.LM("openai/gpt-5", temperature=1.0, max_tokens=16000, cache=False) + # Simulate a pre-existing dual-key state (e.g. produced by some other code path). + lm.kwargs["max_tokens"] = 32000 + + state = lm.dump_state() + + assert state["max_tokens"] == 32000 + assert "max_completion_tokens" not in state + + def test_dump_state(): lm = dspy.LM( model="openai/gpt-4o-mini", From 909bd4b00fe3d3c50ab34daf64a4f3688c5a3b37 Mon Sep 17 00:00:00 2001 From: "detail-app[bot]" Date: Sun, 6 Sep 2026 10:48:21 +0000 Subject: [PATCH 2/2] fix(lm): use target model for token alias check in copy() --- dspy/clients/lm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dspy/clients/lm.py b/dspy/clients/lm.py index 9ffdfa773f..bb9fc0706e 100644 --- a/dspy/clients/lm.py +++ b/dspy/clients/lm.py @@ -409,7 +409,7 @@ def copy(self, **kwargs): the copied LM matches what ``LM.__init__`` would have produced for the same ``max_tokens`` value and the override survives ``dump_state()``. """ - if "max_tokens" in kwargs and _is_openai_reasoning_model(self.model): + if "max_tokens" in kwargs and _is_openai_reasoning_model(kwargs.get("model", self.model)): kwargs["max_completion_tokens"] = kwargs.pop("max_tokens") return super().copy(**kwargs)