feat: expose cumulative token usage on success + budget enforcement

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.
 

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.

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,
         )
 

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