diff --git a/instructor/core/retry.py b/instructor/core/retry.py index 8d586da05..39fefd506 100644 --- a/instructor/core/retry.py +++ b/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) diff --git a/tests/test_token_budget.py b/tests/test_token_budget.py new file mode 100644 index 000000000..06623b8cf --- /dev/null +++ b/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