diff --git a/instructor/v2/core/errors.py b/instructor/v2/core/errors.py index 43bd30be0..e9fc586f8 100644 --- a/instructor/v2/core/errors.py +++ b/instructor/v2/core/errors.py @@ -255,6 +255,39 @@ def __init__( super().__init__(*args, failed_attempts=failed_attempts, **kwargs) +class TokenBudgetExceeded(InstructorError): + """Raised when cumulative token usage exceeds the configured budget. + + This allows users to cap the total cost of retries. When a token_budget + is set and the cumulative tokens across all attempts exceed it, this + exception is raised instead of continuing to retry. + + Attributes: + total_usage: Cumulative token usage at the time of budget breach + budget: The configured token budget that was exceeded + n_attempts: Number of attempts made before budget was exceeded + """ + + def __init__( + self, + *args: Any, + total_usage: Any, + budget: int, + n_attempts: int, + last_completion: Any | None = None, + **kwargs: Any, + ): + self.total_usage = total_usage + self.budget = budget + self.n_attempts = n_attempts + self.last_completion = last_completion + message = ( + f"Token budget exceeded: used {getattr(total_usage, 'total_tokens', '?')} " + f"tokens across {n_attempts} attempts (budget: {budget})" + ) + super().__init__(message, *args, **kwargs) + + class ValidationError(InstructorError): """Exception raised when LLM response validation fails. diff --git a/instructor/v2/core/hooks.py b/instructor/v2/core/hooks.py index 8c1921985..07f217127 100644 --- a/instructor/v2/core/hooks.py +++ b/instructor/v2/core/hooks.py @@ -15,6 +15,7 @@ class HookName(Enum): COMPLETION_RESPONSE = "completion:response" COMPLETION_ERROR = "completion:error" COMPLETION_LAST_ATTEMPT = "completion:last_attempt" + COMPLETION_USAGE = "completion:usage" PARSE_ERROR = "parse:error" @@ -44,6 +45,20 @@ def __call__( ) -> None: ... +class CompletionUsageHandler(Protocol): + """Protocol for completion usage handlers. + + Fired after each API attempt with cumulative token usage. + """ + + def __call__( + self, + usage: Any, + *, + attempt_number: int = ..., + ) -> None: ... + + class ParseErrorHandler(Protocol): """Protocol for parse error handlers.""" @@ -58,6 +73,7 @@ def __call__(self, error: Exception, **kwargs: Any) -> None: ... "completion:response", "completion:error", "completion:last_attempt", + "completion:usage", "parse:error", ], ] @@ -67,6 +83,7 @@ def __call__(self, error: Exception, **kwargs: Any) -> None: ... CompletionKwargsHandler, CompletionResponseHandler, CompletionErrorHandler, + CompletionUsageHandler, ParseErrorHandler, ] @@ -198,6 +215,18 @@ def emit_completion_last_attempt(self, error: Exception, **kwargs: Any) -> None: """ self.emit(HookName.COMPLETION_LAST_ATTEMPT, error, **kwargs) + def emit_completion_usage(self, usage: Any, **kwargs: Any) -> None: + """ + Emit a completion usage event with cumulative token counts. + + Fired after each API attempt with the running total of tokens consumed. + + Args: + usage: Cumulative CompletionUsage object + **kwargs: Optional metadata (attempt_number) + """ + self.emit(HookName.COMPLETION_USAGE, usage, **kwargs) + def emit_parse_error(self, error: Exception, **kwargs: Any) -> None: """ Emit a parse error event. diff --git a/instructor/v2/core/patch.py b/instructor/v2/core/patch.py index b42433bfe..a03a408a5 100644 --- a/instructor/v2/core/patch.py +++ b/instructor/v2/core/patch.py @@ -192,6 +192,7 @@ def new_create_sync( max_retries: int | Retrying = 1, strict: bool = True, hooks: Hooks | None = None, + token_budget: int | None = None, *args: Any, **kwargs: Any, ) -> T_Model: @@ -256,6 +257,7 @@ def new_create_sync( kwargs=new_kwargs, strict=strict, hooks=hooks, + token_budget=token_budget, ) # Store in cache after successful call @@ -301,6 +303,7 @@ async def new_create_async( max_retries: int | AsyncRetrying = 1, strict: bool = True, hooks: Hooks | None = None, + token_budget: int | None = None, *args: Any, **kwargs: Any, ) -> T_Model: @@ -364,6 +367,7 @@ async def new_create_async( args=args, kwargs=new_kwargs, strict=strict, + token_budget=token_budget, hooks=hooks, ) diff --git a/instructor/v2/core/retry.py b/instructor/v2/core/retry.py index f9d9b7db8..6d1deff54 100644 --- a/instructor/v2/core/retry.py +++ b/instructor/v2/core/retry.py @@ -27,6 +27,7 @@ IncompleteOutputException, InstructorRetryException, ResponseParsingError, + TokenBudgetExceeded, ) from instructor.v2.dsl.iterable import IterableBase from instructor.v2.dsl.response_list import ListResponse @@ -69,15 +70,26 @@ def _attempt_metadata( } -def _finalize_parsed_response(parsed: Any, response: Any) -> Any: +def _finalize_parsed_response( + parsed: Any, response: Any, total_usage: Any = None +) -> Any: if isinstance(parsed, IterableBase): parsed = [task for task in parsed.tasks] if isinstance(parsed, AdapterBase): return parsed.content if isinstance(parsed, list) and not isinstance(parsed, ListResponse): - return ListResponse.from_list(parsed, raw_response=response) + result = ListResponse.from_list(parsed, raw_response=response) + if total_usage is not None: + result._total_usage = total_usage # type: ignore[attr-defined] + return result + if isinstance(parsed, ListResponse): + if total_usage is not None: + parsed._total_usage = total_usage # type: ignore[attr-defined] + return parsed if isinstance(parsed, BaseModel): parsed._raw_response = response # type: ignore[attr-defined] + if total_usage is not None: + parsed._total_usage = total_usage # type: ignore[attr-defined] return parsed @@ -121,6 +133,7 @@ def retry_sync_v2( kwargs: dict[str, Any], strict: bool, hooks: Hooks | None = None, + token_budget: int | None = None, ) -> T_Model: """Sync retry logic using v2 registry handlers. @@ -135,12 +148,15 @@ def retry_sync_v2( kwargs: Keyword args for func strict: Strict validation mode hooks: Optional hooks + token_budget: Optional max total tokens across all attempts. If cumulative + usage exceeds this, TokenBudgetExceeded is raised instead of retrying. Returns: Validated Pydantic model instance Raises: InstructorRetryException: If max retries exceeded + TokenBudgetExceeded: If token_budget is set and cumulative usage exceeds it """ if response_model is None: # No structured output, just call the API @@ -207,6 +223,22 @@ def retry_sync_v2( update_total_usage(response=response, total_usage=total_usage) + if hooks: + hooks.emit_completion_usage( + total_usage, attempt_number=attempt_number + ) + + if ( + token_budget is not None + and getattr(total_usage, "total_tokens", 0) > token_budget + ): + raise TokenBudgetExceeded( + total_usage=total_usage, + budget=token_budget, + n_attempts=attempt_number, + last_completion=response, + ) + # Parse response using registry try: stream = kwargs.get("stream", False) @@ -222,7 +254,9 @@ def retry_sync_v2( f"Successfully parsed response on attempt " f"{attempt.retry_state.attempt_number}" ) - return _finalize_parsed_response(parsed, response) + return _finalize_parsed_response( + parsed, response, total_usage=total_usage + ) except IncompleteOutputException: raise @@ -256,7 +290,7 @@ def retry_sync_v2( # Will retry with modified kwargs raise - except IncompleteOutputException: + except (IncompleteOutputException, TokenBudgetExceeded): raise except Exception as e: # Max retries exceeded or non-validation error occurred @@ -366,6 +400,7 @@ async def retry_async_v2( kwargs: dict[str, Any], strict: bool, hooks: Hooks | None = None, + token_budget: int | None = None, ) -> T_Model: """Async retry logic using v2 registry handlers. @@ -380,12 +415,15 @@ async def retry_async_v2( kwargs: Keyword args for func strict: Strict validation mode hooks: Optional hooks + token_budget: Optional max total tokens across all attempts. If cumulative + usage exceeds this, TokenBudgetExceeded is raised instead of retrying. Returns: Validated Pydantic model instance Raises: InstructorRetryException: If max retries exceeded + TokenBudgetExceeded: If token_budget is set and cumulative usage exceeds it """ if response_model is None: # No structured output, just call the API @@ -452,6 +490,22 @@ async def retry_async_v2( update_total_usage(response=response, total_usage=total_usage) + if hooks: + hooks.emit_completion_usage( + total_usage, attempt_number=attempt_number + ) + + if ( + token_budget is not None + and getattr(total_usage, "total_tokens", 0) > token_budget + ): + raise TokenBudgetExceeded( + total_usage=total_usage, + budget=token_budget, + n_attempts=attempt_number, + last_completion=response, + ) + # Parse response using registry try: stream = kwargs.get("stream", False) @@ -467,7 +521,9 @@ async def retry_async_v2( f"Successfully parsed response on attempt " f"{attempt.retry_state.attempt_number}" ) - return _finalize_parsed_response(parsed, response) + return _finalize_parsed_response( + parsed, response, total_usage=total_usage + ) except IncompleteOutputException: raise @@ -501,7 +557,7 @@ async def retry_async_v2( # Will retry with modified kwargs raise - except IncompleteOutputException: + except (IncompleteOutputException, TokenBudgetExceeded): raise except Exception as e: # Max retries exceeded or non-validation error occurred diff --git a/tests/test_token_budget.py b/tests/test_token_budget.py new file mode 100644 index 000000000..72c22f3be --- /dev/null +++ b/tests/test_token_budget.py @@ -0,0 +1,498 @@ +"""Tests for token usage tracking and budget enforcement. + +Verifies that: +1. Total token usage is attached to successful responses (_total_usage) +2. The completion:usage hook fires after each attempt with cumulative usage +3. token_budget parameter raises TokenBudgetExceeded when exceeded +4. _total_usage works for list[Model] / ListResponse shapes +5. Async path has equivalent coverage to sync +""" + +from __future__ import annotations + +import asyncio +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from pydantic import BaseModel + +from instructor.v2.core.errors import TokenBudgetExceeded +from instructor.v2.core.hooks import HookName, Hooks +from instructor.v2.core.mode import Mode +from instructor.v2.core.providers import Provider +from instructor.v2.core.retry import retry_async_v2, retry_sync_v2 +from instructor.v2.dsl.response_list import ListResponse + + +class User(BaseModel): + name: str + age: int + + +def _make_openai_response(content: str, usage_tokens: int = 100): + """Create a mock OpenAI-like response with usage data.""" + from openai.types.completion_usage import ( + CompletionTokensDetails, + CompletionUsage, + PromptTokensDetails, + ) + + response = MagicMock() + response.usage = CompletionUsage( + completion_tokens=usage_tokens // 2, + prompt_tokens=usage_tokens // 2, + total_tokens=usage_tokens, + completion_tokens_details=CompletionTokensDetails( + audio_tokens=0, reasoning_tokens=0 + ), + prompt_tokens_details=PromptTokensDetails(audio_tokens=0, cached_tokens=0), + ) + response.choices = [MagicMock()] + response.choices[0].message = MagicMock() + response.choices[0].message.tool_calls = [MagicMock()] + response.choices[0].message.tool_calls[0].function = MagicMock() + response.choices[0].message.tool_calls[ + 0 + ].function.arguments = '{"name": "Alice", "age": 30}' + return response + + +@patch("instructor.v2.core.retry.mode_registry") +@patch("instructor.v2.core.retry.RegistryValidationMixin") +def test_total_usage_attached_on_success(mock_validation, mock_registry): + """Successful extraction should have _total_usage attached to the result.""" + response = _make_openai_response('{"name": "Alice", "age": 30}', usage_tokens=150) + func = MagicMock(return_value=response) + + handlers = MagicMock() + handlers.response_parser.return_value = User(name="Alice", age=30) + mock_registry.get_handlers.return_value = handlers + + result = retry_sync_v2( + func=func, + response_model=User, + provider=Provider.OPENAI, + mode=Mode.TOOLS, + context=None, + max_retries=3, + args=(), + kwargs={}, + strict=True, + ) + + assert result.name == "Alice" + assert result.age == 30 + assert hasattr(result, "_total_usage") + assert result._total_usage.total_tokens == 150 + + +@patch("instructor.v2.core.retry.mode_registry") +@patch("instructor.v2.core.retry.RegistryValidationMixin") +def test_usage_hook_fires_after_each_attempt(mock_validation, mock_registry): + """completion:usage hook should fire with cumulative usage after each attempt.""" + response = _make_openai_response('{"name": "Bob", "age": 25}', usage_tokens=200) + func = MagicMock(return_value=response) + + handlers = MagicMock() + handlers.response_parser.return_value = User(name="Bob", age=25) + mock_registry.get_handlers.return_value = handlers + + hooks = Hooks() + usage_events: list = [] + + def on_usage(usage, *, attempt_number=0): + usage_events.append({"tokens": usage.total_tokens, "attempt": attempt_number}) + + hooks.on(HookName.COMPLETION_USAGE, on_usage) + + retry_sync_v2( + func=func, + response_model=User, + provider=Provider.OPENAI, + mode=Mode.TOOLS, + context=None, + max_retries=3, + args=(), + kwargs={}, + strict=True, + hooks=hooks, + ) + + assert len(usage_events) == 1 + assert usage_events[0]["tokens"] == 200 + assert usage_events[0]["attempt"] == 1 + + +@patch("instructor.v2.core.retry.mode_registry") +@patch("instructor.v2.core.retry.RegistryValidationMixin") +def test_token_budget_raises_when_exceeded(mock_validation, mock_registry): + """token_budget should raise TokenBudgetExceeded when exceeded.""" + response = _make_openai_response('{"name": "X", "age": 1}', usage_tokens=500) + func = MagicMock(return_value=response) + + handlers = MagicMock() + handlers.response_parser.return_value = User(name="X", age=1) + mock_registry.get_handlers.return_value = handlers + + with pytest.raises(TokenBudgetExceeded) as exc_info: + retry_sync_v2( + func=func, + response_model=User, + provider=Provider.OPENAI, + mode=Mode.TOOLS, + context=None, + max_retries=3, + args=(), + kwargs={}, + strict=True, + token_budget=100, + ) + + assert exc_info.value.budget == 100 + assert exc_info.value.n_attempts == 1 + assert exc_info.value.total_usage.total_tokens == 500 + + +@patch("instructor.v2.core.retry.mode_registry") +@patch("instructor.v2.core.retry.RegistryValidationMixin") +def test_token_budget_none_does_not_limit(mock_validation, mock_registry): + """When token_budget is None, no budget enforcement should occur.""" + response = _make_openai_response('{"name": "Y", "age": 99}', usage_tokens=9999) + func = MagicMock(return_value=response) + + handlers = MagicMock() + handlers.response_parser.return_value = User(name="Y", age=99) + mock_registry.get_handlers.return_value = handlers + + result = retry_sync_v2( + func=func, + response_model=User, + provider=Provider.OPENAI, + mode=Mode.TOOLS, + context=None, + max_retries=3, + args=(), + kwargs={}, + strict=True, + token_budget=None, + ) + + assert result.name == "Y" + assert result._total_usage.total_tokens == 9999 + + +@patch("instructor.v2.core.retry.mode_registry") +@patch("instructor.v2.core.retry.RegistryValidationMixin") +def test_token_budget_allows_under_budget(mock_validation, mock_registry): + """Requests under the token budget should succeed normally.""" + response = _make_openai_response('{"name": "Z", "age": 5}', usage_tokens=50) + func = MagicMock(return_value=response) + + handlers = MagicMock() + handlers.response_parser.return_value = User(name="Z", age=5) + mock_registry.get_handlers.return_value = handlers + + result = retry_sync_v2( + func=func, + response_model=User, + provider=Provider.OPENAI, + mode=Mode.TOOLS, + context=None, + max_retries=3, + args=(), + kwargs={}, + strict=True, + token_budget=1000, + ) + + assert result.name == "Z" + assert result._total_usage.total_tokens == 50 + + +@patch("instructor.v2.core.retry.mode_registry") +@patch("instructor.v2.core.retry.RegistryValidationMixin") +def test_usage_hook_string_registration(mock_validation, mock_registry): + """completion:usage hook should work with string registration.""" + response = _make_openai_response('{"name": "C", "age": 10}', usage_tokens=77) + func = MagicMock(return_value=response) + + handlers = MagicMock() + handlers.response_parser.return_value = User(name="C", age=10) + mock_registry.get_handlers.return_value = handlers + + hooks = Hooks() + called = [] + + def handler(usage, **kwargs): + called.append(usage.total_tokens) + + hooks.on("completion:usage", handler) + + retry_sync_v2( + func=func, + response_model=User, + provider=Provider.OPENAI, + mode=Mode.TOOLS, + context=None, + max_retries=1, + args=(), + kwargs={}, + strict=True, + hooks=hooks, + ) + + assert called == [77] + + +# --- list[Model] / ListResponse tests --- + + +@patch("instructor.v2.core.retry.mode_registry") +@patch("instructor.v2.core.retry.RegistryValidationMixin") +def test_total_usage_attached_to_list_response(mock_validation, mock_registry): + """_total_usage should be attached to ListResponse results.""" + response = _make_openai_response('[{"name": "A", "age": 1}]', usage_tokens=200) + func = MagicMock(return_value=response) + + users = [User(name="A", age=1), User(name="B", age=2)] + + handlers = MagicMock() + handlers.response_parser.return_value = users + mock_registry.get_handlers.return_value = handlers + + result = retry_sync_v2( + func=func, + response_model=User, + provider=Provider.OPENAI, + mode=Mode.TOOLS, + context=None, + max_retries=3, + args=(), + kwargs={}, + strict=True, + ) + + assert isinstance(result, ListResponse) + assert len(result) == 2 + assert hasattr(result, "_total_usage") + assert result._total_usage.total_tokens == 200 + + +@patch("instructor.v2.core.retry.mode_registry") +@patch("instructor.v2.core.retry.RegistryValidationMixin") +def test_total_usage_on_list_response_with_budget(mock_validation, mock_registry): + """token_budget should work correctly with ListResponse results.""" + response = _make_openai_response('[{"name": "X", "age": 5}]', usage_tokens=50) + func = MagicMock(return_value=response) + + handlers = MagicMock() + handlers.response_parser.return_value = [User(name="X", age=5)] + mock_registry.get_handlers.return_value = handlers + + result = retry_sync_v2( + func=func, + response_model=User, + provider=Provider.OPENAI, + mode=Mode.TOOLS, + context=None, + max_retries=3, + args=(), + kwargs={}, + strict=True, + token_budget=1000, + ) + + assert isinstance(result, ListResponse) + assert result._total_usage.total_tokens == 50 + + +@patch("instructor.v2.core.retry.mode_registry") +@patch("instructor.v2.core.retry.RegistryValidationMixin") +def test_token_budget_raises_for_list_response(mock_validation, mock_registry): + """TokenBudgetExceeded should fire for list responses too.""" + response = _make_openai_response('[{"name": "X", "age": 5}]', usage_tokens=500) + func = MagicMock(return_value=response) + + handlers = MagicMock() + handlers.response_parser.return_value = [User(name="X", age=5)] + mock_registry.get_handlers.return_value = handlers + + with pytest.raises(TokenBudgetExceeded) as exc_info: + retry_sync_v2( + func=func, + response_model=User, + provider=Provider.OPENAI, + mode=Mode.TOOLS, + context=None, + max_retries=3, + args=(), + kwargs={}, + strict=True, + token_budget=100, + ) + + assert exc_info.value.total_usage.total_tokens == 500 + + +# --- Async tests --- + + +@patch("instructor.v2.core.retry.mode_registry") +@patch("instructor.v2.core.retry.RegistryValidationMixin") +def test_async_total_usage_attached_on_success(mock_validation, mock_registry): + """Async: _total_usage should be attached to successful results.""" + response = _make_openai_response('{"name": "Alice", "age": 30}', usage_tokens=150) + func = AsyncMock(return_value=response) + + handlers = MagicMock() + handlers.response_parser.return_value = User(name="Alice", age=30) + mock_registry.get_handlers.return_value = handlers + + result = asyncio.run( + retry_async_v2( + func=func, + response_model=User, + provider=Provider.OPENAI, + mode=Mode.TOOLS, + context=None, + max_retries=3, + args=(), + kwargs={}, + strict=True, + ) + ) + + assert result.name == "Alice" + assert hasattr(result, "_total_usage") + assert result._total_usage.total_tokens == 150 + + +@patch("instructor.v2.core.retry.mode_registry") +@patch("instructor.v2.core.retry.RegistryValidationMixin") +def test_async_usage_hook_fires(mock_validation, mock_registry): + """Async: completion:usage hook should fire with cumulative usage.""" + response = _make_openai_response('{"name": "Bob", "age": 25}', usage_tokens=200) + func = AsyncMock(return_value=response) + + handlers = MagicMock() + handlers.response_parser.return_value = User(name="Bob", age=25) + mock_registry.get_handlers.return_value = handlers + + hooks = Hooks() + usage_events: list = [] + + def on_usage(usage, *, attempt_number=0): + usage_events.append({"tokens": usage.total_tokens, "attempt": attempt_number}) + + hooks.on(HookName.COMPLETION_USAGE, on_usage) + + asyncio.run( + retry_async_v2( + func=func, + response_model=User, + provider=Provider.OPENAI, + mode=Mode.TOOLS, + context=None, + max_retries=3, + args=(), + kwargs={}, + strict=True, + hooks=hooks, + ) + ) + + assert len(usage_events) == 1 + assert usage_events[0]["tokens"] == 200 + assert usage_events[0]["attempt"] == 1 + + +@patch("instructor.v2.core.retry.mode_registry") +@patch("instructor.v2.core.retry.RegistryValidationMixin") +def test_async_token_budget_raises_when_exceeded(mock_validation, mock_registry): + """Async: token_budget should raise TokenBudgetExceeded.""" + response = _make_openai_response('{"name": "X", "age": 1}', usage_tokens=500) + func = AsyncMock(return_value=response) + + handlers = MagicMock() + handlers.response_parser.return_value = User(name="X", age=1) + mock_registry.get_handlers.return_value = handlers + + with pytest.raises(TokenBudgetExceeded) as exc_info: + asyncio.run( + retry_async_v2( + func=func, + response_model=User, + provider=Provider.OPENAI, + mode=Mode.TOOLS, + context=None, + max_retries=3, + args=(), + kwargs={}, + strict=True, + token_budget=100, + ) + ) + + assert exc_info.value.budget == 100 + assert exc_info.value.n_attempts == 1 + assert exc_info.value.total_usage.total_tokens == 500 + + +@patch("instructor.v2.core.retry.mode_registry") +@patch("instructor.v2.core.retry.RegistryValidationMixin") +def test_async_token_budget_allows_under_budget(mock_validation, mock_registry): + """Async: requests under budget should succeed normally.""" + response = _make_openai_response('{"name": "Z", "age": 5}', usage_tokens=50) + func = AsyncMock(return_value=response) + + handlers = MagicMock() + handlers.response_parser.return_value = User(name="Z", age=5) + mock_registry.get_handlers.return_value = handlers + + result = asyncio.run( + retry_async_v2( + func=func, + response_model=User, + provider=Provider.OPENAI, + mode=Mode.TOOLS, + context=None, + max_retries=3, + args=(), + kwargs={}, + strict=True, + token_budget=1000, + ) + ) + + assert result.name == "Z" + assert result._total_usage.total_tokens == 50 + + +@patch("instructor.v2.core.retry.mode_registry") +@patch("instructor.v2.core.retry.RegistryValidationMixin") +def test_async_list_response_has_total_usage(mock_validation, mock_registry): + """Async: list responses should also have _total_usage attached.""" + response = _make_openai_response('[{"name": "A", "age": 1}]', usage_tokens=300) + func = AsyncMock(return_value=response) + + handlers = MagicMock() + handlers.response_parser.return_value = [User(name="A", age=1)] + mock_registry.get_handlers.return_value = handlers + + result = asyncio.run( + retry_async_v2( + func=func, + response_model=User, + provider=Provider.OPENAI, + mode=Mode.TOOLS, + context=None, + max_retries=3, + args=(), + kwargs={}, + strict=True, + ) + ) + + assert isinstance(result, ListResponse) + assert hasattr(result, "_total_usage") + assert result._total_usage.total_tokens == 300