From 998820d85147a838222dbc7710ecaf3b71a2c064 Mon Sep 17 00:00:00 2001 From: mimran-khan Date: Thu, 25 Jun 2026 01:05:52 +0530 Subject: [PATCH 1/2] feat: expose cumulative token usage on success + budget enforcement Previously, total_usage (accumulated token counts across retry attempts) was only available when retries FAILED via InstructorRetryException. On successful extraction, this data was computed but discarded. Changes: - Attach _total_usage to successful parsed responses (like _raw_response) - Add completion:usage hook that fires after each attempt with cumulative token counts, enabling observability integration - Add token_budget parameter to create() that raises TokenBudgetExceeded if cumulative tokens across all attempts exceed the configured limit - TokenBudgetExceeded exception includes usage data, budget, and attempt count This prevents runaway retry costs in production and gives users visibility into how many tokens their retries actually consume when things work. Refs #2391 --- instructor/v2/core/errors.py | 33 +++++ instructor/v2/core/hooks.py | 29 +++++ instructor/v2/core/patch.py | 4 + instructor/v2/core/retry.py | 59 ++++++++- tests/test_token_budget.py | 240 +++++++++++++++++++++++++++++++++++ 5 files changed, 360 insertions(+), 5 deletions(-) create mode 100644 tests/test_token_budget.py 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..7eb0ef3db 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,7 +70,9 @@ 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): @@ -78,6 +81,8 @@ def _finalize_parsed_response(parsed: Any, response: Any) -> Any: return ListResponse.from_list(parsed, raw_response=response) 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 +126,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 +141,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 +216,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 +247,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 +283,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 +393,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 +408,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 +483,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 +514,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 +550,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..b1788d596 --- /dev/null +++ b/tests/test_token_budget.py @@ -0,0 +1,240 @@ +"""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 +""" + +from __future__ import annotations + +from unittest.mock import 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_sync_v2 + + +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 cumulative tokens exceed it.""" + 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] From fe057cdab63c51db35e5ad5c5faed7b3ada7e290 Mon Sep 17 00:00:00 2001 From: mimran-khan Date: Mon, 29 Jun 2026 19:52:30 +0530 Subject: [PATCH 2/2] feat: attach _total_usage to ListResponse + add async test coverage Addresses reviewer feedback: - _total_usage is now attached to ListResponse results (list[Model]), not just single BaseModel responses - Added 5 async tests mirroring the sync coverage (usage attachment, hook firing, budget enforcement, under-budget success, list response) - Added 3 ListResponse-specific tests (sync) for budget enforcement and usage attachment on list results AdapterBase (primitive return types like str/int) cannot carry attributes, so _total_usage is only available on BaseModel and ListResponse shapes. --- instructor/v2/core/retry.py | 9 +- tests/test_token_budget.py | 270 +++++++++++++++++++++++++++++++++++- 2 files changed, 272 insertions(+), 7 deletions(-) diff --git a/instructor/v2/core/retry.py b/instructor/v2/core/retry.py index 7eb0ef3db..6d1deff54 100644 --- a/instructor/v2/core/retry.py +++ b/instructor/v2/core/retry.py @@ -78,7 +78,14 @@ def _finalize_parsed_response( 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: diff --git a/tests/test_token_budget.py b/tests/test_token_budget.py index b1788d596..72c22f3be 100644 --- a/tests/test_token_budget.py +++ b/tests/test_token_budget.py @@ -4,11 +4,14 @@ 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 -from unittest.mock import MagicMock, patch +import asyncio +from unittest.mock import AsyncMock, MagicMock, patch import pytest from pydantic import BaseModel @@ -17,7 +20,8 @@ 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_sync_v2 +from instructor.v2.core.retry import retry_async_v2, retry_sync_v2 +from instructor.v2.dsl.response_list import ListResponse class User(BaseModel): @@ -47,9 +51,9 @@ def _make_openai_response(content: str, usage_tokens: int = 100): 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}' - ) + response.choices[0].message.tool_calls[ + 0 + ].function.arguments = '{"name": "Alice", "age": 30}' return response @@ -122,7 +126,7 @@ def on_usage(usage, *, attempt_number=0): @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 cumulative tokens exceed it.""" + """token_budget should raise TokenBudgetExceeded when exceeded.""" response = _make_openai_response('{"name": "X", "age": 1}', usage_tokens=500) func = MagicMock(return_value=response) @@ -238,3 +242,257 @@ def handler(usage, **kwargs): ) 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