Skip to content

fix(dspy): avoid KeyError in truncation warning for reasoning models - #88

Open
detail-app[bot] wants to merge 1 commit into
mainfrom
detail/bug-fix/fix-dspy-avoid-keyerror-in-truncation-warning-for-81603e
Open

fix(dspy): avoid KeyError in truncation warning for reasoning models#88
detail-app[bot] wants to merge 1 commit into
mainfrom
detail/bug-fix/fix-dspy-avoid-keyerror-in-truncation-warning-for-81603e

Conversation

@detail-app

@detail-app detail-app Bot commented Sep 6, 2026

Copy link
Copy Markdown

Warning

GitHub issue creation failed

Detail attempted to publish this bug to GitHub, but the issue could not be created. This fix PR was created without that issue, and missing tracker references are shown as Unknown issue.

You can review and merge this PR normally. Please review your tracker integration settings before the next publish run.

Detail bug report: View on Detail

📝 Changes Description

Closes Unknown issue

Bug: LM._check_truncation (dspy/clients/lm.py) interpolates the configured token limit into its truncation warning via self.kwargs['max_tokens']. OpenAI reasoning models (openai/o3, o1, gpt-5, etc.) store that limit under max_completion_tokens instead — an asserted invariant ("max_tokens" not in lm.kwargs). So a truncated chat completion (finish_reason == "length") raised KeyError: 'max_tokens' out of both LM.forward and LM.aforward, turning a successful provider response into an unhandled crash. This was missed when _check_truncation was added without reconciling it against the reasoning-model kwargs shape.

Fix: Read the token limit defensively from whichever key is present, then interpolate the local variable:

max_tokens = self.kwargs.get("max_tokens", self.kwargs.get("max_completion_tokens"))

This is the minimal idiomatic fix — it preserves the reasoning-model kwargs invariant (no change to self.kwargs) and covers both reasoning (max_completion_tokens) and non-reasoning (max_tokens) chat paths. forward and aforward share the method, so both are fixed.

Testing:

  • Added 3 unit tests in tests/clients/test_lm.py covering the previously-untested _check_truncation: a reasoning-model truncated response warns and returns the partial result (the exact bug path), a non-reasoning truncated response still warns (the other branch of the defensive read), and a non-truncated response emits no warning (guards the precondition).
  • Confirmed the reasoning test reproduces the exact KeyError: 'max_tokens' at dspy/clients/lm.py:440 when the fix is reverted, and passes with the fix in place.
  • Routine checks pass: pytest tests/clients/test_lm.py (90 passed), the full tests/clients/ run (157 passed, 22 live-only skipped), ruff check clean, and pre-commit run --files passed.
  • Also confirmed end-to-end via dspy.Predict with a mocked reasoning-model truncated response — the adapter returns partial output and logs the warning instead of raising.
  • Could not run live verification against the real OpenAI Chat Completions endpoint: OPENAI_API_KEY is not set in this environment, so the provider call fails authentication before _check_truncation is reached. The mocked tests cover the same code path up to the provider boundary.

✅ Contributor Checklist

  • Pre-Commit checks are passing (locally and remotely)
  • Title of your PR / MR corresponds to the required format
  • Commit message follows required format {label}(dspy): {message}

⚠️ Warnings

  • ruff format --check reports drift in a few pre-existing, untouched lines of dspy/clients/lm.py and tests/clients/test_lm.py (identical at HEAD); the changed lines are format-compliant, so this PR does not introduce new format drift.
  • Live reasoning-model truncation could not be exercised without an OpenAI API key (see Testing).

AI disclosure: this fix was produced with assistance from Detail (an AI coding assistant).


Automatic Fixes PRs can be configured here.

@greptile-apps

greptile-apps Bot commented Sep 6, 2026

Copy link
Copy Markdown

Greptile Summary

This PR prevents truncated responses from reasoning models from raising KeyError while constructing the truncation warning.

  • Reads the configured limit from either max_tokens or max_completion_tokens.
  • Preserves existing behavior for non-reasoning chat models.
  • Adds regression coverage for reasoning-model truncation, ordinary chat truncation, and non-truncated responses.

Confidence Score: 5/5

The PR appears safe to merge and fixes the reported reasoning-model truncation crash without changing provider-request behavior.

Reasoning-model construction stores the limit as max_completion_tokens, while ordinary chat construction stores it as max_tokens; the fallback handles both configurations, and the shared truncation checker covers synchronous and asynchronous calls.

Important Files Changed

Filename Overview
dspy/clients/lm.py Safely retrieves either supported token-limit key before logging a truncation warning.
tests/clients/test_lm.py Adds representative regression tests covering both token-limit key shapes and the non-truncated path.

Reviews (1): Last reviewed commit: "fix(dspy): avoid KeyError in truncation ..." | Re-trigger Greptile

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant