feat(core): add token_budget parameter to limit retry token usage

instructor/core/retry.py

@@ -49,6 +49,8 @@ def initialize_retrying(
     max_retries: int | Retrying | AsyncRetrying,
     is_async: bool,
     timeout: float | None = None,
+    token_budget: int | None = None,
+    usage_container: list | None = None,
 ):
     """
     Initialize the retrying mechanism based on the type (synchronous or asynchronous).
@@ -57,6 +59,8 @@ def initialize_retrying(
         max_retries (int | Retrying | AsyncRetrying): Maximum number of retries or a retrying object.
         is_async (bool): Flag indicating if the retrying is asynchronous.
         timeout (float | None): Optional timeout in seconds to limit total retry duration.
+        token_budget (int | None): Optional max total tokens before stopping retries.
+        usage_container (list | None): Mutable container holding the running CompletionUsage object.
 
     Returns:
         Retrying | AsyncRetrying: Configured retrying object.
@@ -70,6 +74,19 @@ def initialize_retrying(
             # Add global timeout: stop after timeout seconds total
             stop_conditions.append(stop_after_delay(timeout))
 
+        if token_budget is not None and usage_container is not None:
+
+            def _stop_on_token_budget(retry_state: Any) -> bool:  # noqa: ARG001
+                usage = usage_container[0]
+                total = getattr(usage, "total_tokens", None)
+                if total is None:
+                    total = getattr(usage, "input_tokens", 0) + getattr(
+                        usage, "output_tokens", 0
+                    )
+                return total >= token_budget
+
+            stop_conditions.append(_stop_on_token_budget)
+
         # Combine stop conditions with OR logic (stop if ANY condition is met)
         stop_condition = stop_conditions[0]
         for condition in stop_conditions[1:]:
@@ -178,9 +195,17 @@ def retry_sync(
     """
     hooks = hooks or Hooks()
     total_usage = initialize_usage(mode)
-    # Extract timeout from kwargs if available (for global timeout across retries)
+    # Extract timeout and token_budget from kwargs if available
     timeout = kwargs.get("timeout")
-    max_retries = initialize_retrying(max_retries, is_async=False, timeout=timeout)
+    token_budget = kwargs.pop("token_budget", None)
+    usage_container = [total_usage]
+    max_retries = initialize_retrying(
+        max_retries,
+        is_async=False,
+        timeout=timeout,
+        token_budget=token_budget,
+        usage_container=usage_container,
+    )
 
     # Pre-extract stream flag to avoid repeated lookup
     stream = kwargs.get("stream", False)
@@ -356,9 +381,17 @@ async def retry_async(
     """
     hooks = hooks or Hooks()
     total_usage = initialize_usage(mode)
-    # Extract timeout from kwargs if available (for global timeout across retries)
+    # Extract timeout and token_budget from kwargs if available
     timeout = kwargs.get("timeout")
-    max_retries = initialize_retrying(max_retries, is_async=True, timeout=timeout)
+    token_budget = kwargs.pop("token_budget", None)
+    usage_container = [total_usage]
+    max_retries = initialize_retrying(
+        max_retries,
+        is_async=True,
+        timeout=timeout,
+        token_budget=token_budget,
+        usage_container=usage_container,
+    )
 
     # Pre-extract stream flag to avoid repeated lookup
     stream = kwargs.get("stream", False)

tests/test_token_budget.py

@@ -0,0 +1,120 @@
+"""
+Test that token_budget stops retries when cumulative token usage exceeds the budget.
+
+This tests the mitigation for retry amplification (issue #2056).
+"""
+
+from unittest.mock import Mock
+
+import pytest
+from openai.types.chat import ChatCompletion, ChatCompletionMessage
+from openai.types.chat.chat_completion import Choice
+from openai.types.completion_usage import CompletionUsage
+from pydantic import BaseModel, field_validator
+
+import instructor
+from instructor.core.exceptions import InstructorRetryException
+from instructor.mode import Mode
+
+
+class StrictAge(BaseModel):
+    name: str
+    age: int
+
+    @field_validator("age")
+    @classmethod
+    def age_must_be_positive(cls, v: int) -> int:
+        if v < 0:
+            raise ValueError("age must be positive")
+        return v
+
+
+def _make_completion(content: str, usage_tokens: int) -> ChatCompletion:
+    return ChatCompletion(
+        id="test",
+        model="gpt-4",
+        object="chat.completion",
+        created=0,
+        choices=[
+            Choice(
+                index=0,
+                message=ChatCompletionMessage(role="assistant", content=content),
+                finish_reason="stop",
+            )
+        ],
+        usage=CompletionUsage(
+            prompt_tokens=usage_tokens // 2,
+            completion_tokens=usage_tokens // 2,
+            total_tokens=usage_tokens,
+        ),
+    )
+
+
+def test_token_budget_stops_retries():
+    """Retries stop when cumulative token usage exceeds token_budget."""
+    # Each call uses 500 tokens, budget is 800 -- should stop after 2 attempts
+    bad_response = _make_completion('{"name": "Alice", "age": -1}', usage_tokens=500)
+
+    mock_client = Mock()
+    mock_client.chat = Mock()
+    mock_client.chat.completions = Mock()
+    mock_client.chat.completions.create = Mock(return_value=bad_response)
+
+    client = instructor.patch(mock_client, mode=Mode.JSON)
+
+    with pytest.raises(InstructorRetryException) as exc_info:
+        client.chat.completions.create(
+            model="gpt-4",
+            response_model=StrictAge,
+            messages=[{"role": "user", "content": "give me a user"}],
+            max_retries=10,
+            token_budget=800,
+        )
+
+    # Should have stopped well before 10 retries due to budget
+    assert exc_info.value.n_attempts <= 3
+
+
+def test_token_budget_not_set_retries_normally():
+    """Without token_budget, all retries are exhausted."""
+    bad_response = _make_completion('{"name": "Alice", "age": -1}', usage_tokens=500)
+
+    mock_client = Mock()
+    mock_client.chat = Mock()
+    mock_client.chat.completions = Mock()
+    mock_client.chat.completions.create = Mock(return_value=bad_response)
+
+    client = instructor.patch(mock_client, mode=Mode.JSON)
+
+    with pytest.raises(InstructorRetryException) as exc_info:
+        client.chat.completions.create(
+            model="gpt-4",
+            response_model=StrictAge,
+            messages=[{"role": "user", "content": "give me a user"}],
+            max_retries=5,
+        )
+
+    assert exc_info.value.n_attempts == 5
+
+
+def test_token_budget_success_before_limit():
+    """If validation succeeds before budget is hit, result is returned normally."""
+    good_response = _make_completion('{"name": "Alice", "age": 30}', usage_tokens=200)
+
+    mock_client = Mock()
+    mock_client.chat = Mock()
+    mock_client.chat.completions = Mock()
+    mock_client.chat.completions.create = Mock(return_value=good_response)
+
+    client = instructor.patch(mock_client, mode=Mode.JSON)
+
+    result = client.chat.completions.create(
+        model="gpt-4",
+        response_model=StrictAge,
+        messages=[{"role": "user", "content": "give me a user"}],
+        max_retries=5,
+        token_budget=1000,
+    )
+
+    assert result.name == "Alice"
+    assert result.age == 30