Files
wehub-resource-sync 97e91a83f3
Ruff / Ruff (push) Has been cancelled
Test / Core Tests (push) Has been cancelled
Test / Offline Coverage Tests (Python 3.10) (push) Has been cancelled
Test / Offline Coverage Tests (Python 3.11) (push) Has been cancelled
Test / Offline Coverage Tests (Python 3.12) (push) Has been cancelled
Test / Offline Coverage Tests (Python 3.13) (push) Has been cancelled
Test / Offline Coverage Tests (Python 3.9) (push) Has been cancelled
Test / Full Coverage (Python 3.11) (push) Has been cancelled
Test / Core Provider Tests (OpenAI) (push) Has been cancelled
Test / Core Provider Tests (Anthropic) (push) Has been cancelled
Test / Core Provider Tests (Google) (push) Has been cancelled
Test / Core Provider Tests (Other) (push) Has been cancelled
Test / Anthropic Tests (push) Has been cancelled
Test / Gemini Tests (push) Has been cancelled
Test / Google GenAI Tests (push) Has been cancelled
Test / Vertex AI Tests (push) Has been cancelled
Test / OpenAI Tests (push) Has been cancelled
Test / Writer Tests (push) Has been cancelled
Test / Auto Client Tests (push) Has been cancelled
ty / type-check (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 13:36:38 +08:00

545 lines
18 KiB
Python

"""v2 retry mechanism using registry handlers.
Custom retry logic for v2 that uses registry's reask and response_parser
instead of v1's process_response.
"""
from __future__ import annotations
import json
import logging
from typing import TYPE_CHECKING, Any, TypeVar
from pydantic import BaseModel, ValidationError
from tenacity import (
AsyncRetrying,
Retrying,
retry_if_exception_type,
stop_after_attempt,
stop_after_delay,
)
from instructor.v2.core.mode import Mode
from instructor.v2.core.providers import Provider
from instructor.v2.core.errors import (
AsyncValidationError,
FailedAttempt,
IncompleteOutputException,
InstructorRetryException,
ResponseParsingError,
)
from instructor.v2.dsl.iterable import IterableBase
from instructor.v2.dsl.response_list import ListResponse
from instructor.v2.dsl.simple_type import AdapterBase
from instructor.v2.core.messages import extract_messages
from instructor.v2.core.usage import update_total_usage
from instructor.v2.core.exceptions import RegistryValidationMixin
from instructor.v2.core.registry import mode_registry
if TYPE_CHECKING:
from collections.abc import Awaitable, Callable
from instructor.v2.core.hooks import Hooks
logger = logging.getLogger("instructor.v2.retry")
T_Model = TypeVar("T_Model", bound=BaseModel)
_RETRYABLE_PARSE_ERRORS = (
ValidationError,
json.JSONDecodeError,
AsyncValidationError,
ResponseParsingError,
)
def _max_attempts(max_retries: int | Retrying | AsyncRetrying) -> int | None:
return max(max_retries, 0) + 1 if isinstance(max_retries, int) else None
def _attempt_metadata(
*,
attempt_number: int,
max_attempts: int | None,
is_last_attempt: bool,
) -> dict[str, Any]:
return {
"attempt_number": attempt_number,
"max_attempts": max_attempts,
"is_last_attempt": is_last_attempt,
}
def _finalize_parsed_response(parsed: Any, response: Any) -> 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)
if isinstance(parsed, BaseModel):
parsed._raw_response = response # type: ignore[attr-defined]
return parsed
def _initialize_usage(provider: Provider | Mode) -> Any:
from openai.types.completion_usage import (
CompletionTokensDetails,
CompletionUsage,
PromptTokensDetails,
)
total_usage: Any = CompletionUsage(
completion_tokens=0,
prompt_tokens=0,
total_tokens=0,
completion_tokens_details=CompletionTokensDetails(
audio_tokens=0, reasoning_tokens=0
),
prompt_tokens_details=PromptTokensDetails(audio_tokens=0, cached_tokens=0),
)
anthropic_modes = {
Mode.ANTHROPIC_TOOLS,
Mode.ANTHROPIC_REASONING_TOOLS,
Mode.ANTHROPIC_JSON,
Mode.ANTHROPIC_PARALLEL_TOOLS,
}
if provider is Provider.ANTHROPIC or provider in anthropic_modes:
from instructor.v2.providers.anthropic.usage import initialize_usage
total_usage = initialize_usage()
return total_usage
def retry_sync_v2(
func: Callable[..., Any],
response_model: type[T_Model] | None,
provider: Provider,
mode: Mode,
context: dict[str, Any] | None,
max_retries: int | Retrying,
args: tuple[Any, ...],
kwargs: dict[str, Any],
strict: bool,
hooks: Hooks | None = None,
) -> T_Model:
"""Sync retry logic using v2 registry handlers.
Args:
func: API function to call
response_model: Pydantic model to extract
provider: Provider enum
mode: Mode enum
context: Validation context
max_retries: Maximum retries after the initial attempt, or Retrying instance
args: Positional args for func
kwargs: Keyword args for func
strict: Strict validation mode
hooks: Optional hooks
Returns:
Validated Pydantic model instance
Raises:
InstructorRetryException: If max retries exceeded
"""
if response_model is None:
# No structured output, just call the API
return func(*args, **kwargs)
# Validate and get handlers from registry
RegistryValidationMixin.validate_mode_registration(provider, mode)
handlers = mode_registry.get_handlers(provider, mode)
# Setup retrying
if isinstance(max_retries, int):
stop_condition = stop_after_attempt(max(max_retries, 0) + 1)
timeout = kwargs.get("timeout")
if isinstance(timeout, (int, float)):
stop_condition = stop_condition | stop_after_delay(timeout)
max_retries_instance = Retrying(
stop=stop_condition,
retry=retry_if_exception_type(_RETRYABLE_PARSE_ERRORS),
reraise=True,
)
else:
max_retries_instance = max_retries
max_attempts = _max_attempts(max_retries)
failed_attempts: list[FailedAttempt] = []
last_exception: Exception | None = None
last_attempt_number = 0
total_usage = _initialize_usage(provider)
try:
for attempt in max_retries_instance:
with attempt:
attempt_number = attempt.retry_state.attempt_number
last_attempt_number = attempt_number
# Call API
if hooks:
hooks.emit_completion_arguments(**kwargs)
try:
response = func(*args, **kwargs)
except IncompleteOutputException:
raise
except Exception as e:
logger.error(f"API call failed on attempt {attempt_number}: {e}")
if hooks:
hooks.emit_completion_error(
e,
**_attempt_metadata(
attempt_number=attempt_number,
max_attempts=max_attempts,
is_last_attempt=(
not isinstance(e, ValidationError)
or (
max_attempts is not None
and attempt_number >= max_attempts
)
),
),
)
raise
if hooks:
hooks.emit_completion_response(response)
update_total_usage(response=response, total_usage=total_usage)
# Parse response using registry
try:
stream = kwargs.get("stream", False)
parsed = handlers.response_parser(
response=response,
response_model=response_model,
validation_context=context,
strict=strict,
stream=stream,
is_async=False,
)
logger.debug(
f"Successfully parsed response on attempt "
f"{attempt.retry_state.attempt_number}"
)
return _finalize_parsed_response(parsed, response)
except IncompleteOutputException:
raise
except _RETRYABLE_PARSE_ERRORS as e:
logger.debug(f"Validation error on attempt {attempt_number}: {e}")
failed_attempts.append(
FailedAttempt(
attempt_number=attempt_number,
exception=e,
completion=response,
)
)
last_exception = e
if hooks:
hooks.emit_parse_error(
e,
**_attempt_metadata(
attempt_number=attempt_number,
max_attempts=max_attempts,
is_last_attempt=max_attempts == attempt_number,
),
)
# Prepare reask using registry
kwargs = handlers.reask_handler(
kwargs=kwargs,
response=response,
exception=e,
)
# Will retry with modified kwargs
raise
except IncompleteOutputException:
raise
except Exception as e:
# Max retries exceeded or non-validation error occurred
last_exception = e
logger.error(
f"Max retries exceeded. Total attempts: {last_attempt_number}, "
f"Last error: {last_exception}"
)
if hooks:
hooks.emit_completion_last_attempt(
last_exception,
**_attempt_metadata(
attempt_number=last_attempt_number or len(failed_attempts),
max_attempts=max_attempts,
is_last_attempt=True,
),
)
raise InstructorRetryException(
str(last_exception),
last_completion=failed_attempts[-1].completion if failed_attempts else None,
n_attempts=last_attempt_number,
total_usage=total_usage,
messages=extract_messages(kwargs),
create_kwargs=kwargs,
failed_attempts=failed_attempts,
) from last_exception
# Should never reach here
logger.error("Unexpected code path in retry_sync_v2")
raise InstructorRetryException(
str(last_exception) if last_exception else "Unknown error",
last_completion=failed_attempts[-1].completion if failed_attempts else None,
n_attempts=last_attempt_number,
total_usage=total_usage,
messages=extract_messages(kwargs),
create_kwargs=kwargs,
failed_attempts=failed_attempts,
)
def retry_sync(
func: Callable[..., Any],
response_model: type[T_Model] | None,
args: Any,
kwargs: Any,
context: dict[str, Any] | None = None,
max_retries: int | Retrying = 1,
strict: bool | None = None,
mode: Mode = Mode.TOOLS,
provider: Provider = Provider.OPENAI,
hooks: Hooks | None = None,
) -> T_Model | None:
"""Compatibility wrapper for the public retry API."""
strict_value = True if strict is None else strict
return retry_sync_v2(
func=func,
response_model=response_model,
provider=provider,
mode=mode,
context=context,
max_retries=max_retries,
args=tuple(args) if isinstance(args, tuple) else args,
kwargs=dict(kwargs),
strict=strict_value,
hooks=hooks,
)
async def retry_async(
func: Callable[..., Any],
response_model: type[T_Model] | None,
args: Any,
kwargs: Any,
context: dict[str, Any] | None = None,
max_retries: int | AsyncRetrying = 1,
strict: bool | None = None,
mode: Mode = Mode.TOOLS,
provider: Provider = Provider.OPENAI,
hooks: Hooks | None = None,
) -> T_Model | None:
"""Compatibility wrapper for the public retry API."""
strict_value = True if strict is None else strict
return await retry_async_v2(
func=func,
response_model=response_model,
provider=provider,
mode=mode,
context=context,
max_retries=max_retries,
args=tuple(args) if isinstance(args, tuple) else args,
kwargs=dict(kwargs),
strict=strict_value,
hooks=hooks,
)
async def retry_async_v2(
func: Callable[..., Awaitable[Any]],
response_model: type[T_Model] | None,
provider: Provider,
mode: Mode,
context: dict[str, Any] | None,
max_retries: int | AsyncRetrying,
args: tuple[Any, ...],
kwargs: dict[str, Any],
strict: bool,
hooks: Hooks | None = None,
) -> T_Model:
"""Async retry logic using v2 registry handlers.
Args:
func: Async API function to call
response_model: Pydantic model to extract
provider: Provider enum
mode: Mode enum
context: Validation context
max_retries: Maximum retries after the initial attempt, or AsyncRetrying instance
args: Positional args for func
kwargs: Keyword args for func
strict: Strict validation mode
hooks: Optional hooks
Returns:
Validated Pydantic model instance
Raises:
InstructorRetryException: If max retries exceeded
"""
if response_model is None:
# No structured output, just call the API
return await func(*args, **kwargs)
# Validate and get handlers from registry
RegistryValidationMixin.validate_mode_registration(provider, mode)
handlers = mode_registry.get_handlers(provider, mode)
# Setup retrying
if isinstance(max_retries, int):
stop_condition = stop_after_attempt(max(max_retries, 0) + 1)
timeout = kwargs.get("timeout")
if isinstance(timeout, (int, float)):
stop_condition = stop_condition | stop_after_delay(timeout)
max_retries_instance = AsyncRetrying(
stop=stop_condition,
retry=retry_if_exception_type(_RETRYABLE_PARSE_ERRORS),
reraise=True,
)
else:
max_retries_instance = max_retries
max_attempts = _max_attempts(max_retries)
failed_attempts: list[FailedAttempt] = []
last_exception: Exception | None = None
last_attempt_number = 0
total_usage = _initialize_usage(provider)
try:
async for attempt in max_retries_instance:
with attempt:
attempt_number = attempt.retry_state.attempt_number
last_attempt_number = attempt_number
# Call API
if hooks:
hooks.emit_completion_arguments(**kwargs)
try:
response = await func(*args, **kwargs)
except IncompleteOutputException:
raise
except Exception as e:
logger.error(f"API call failed on attempt {attempt_number}: {e}")
if hooks:
hooks.emit_completion_error(
e,
**_attempt_metadata(
attempt_number=attempt_number,
max_attempts=max_attempts,
is_last_attempt=(
not isinstance(e, ValidationError)
or (
max_attempts is not None
and attempt_number >= max_attempts
)
),
),
)
raise
if hooks:
hooks.emit_completion_response(response)
update_total_usage(response=response, total_usage=total_usage)
# Parse response using registry
try:
stream = kwargs.get("stream", False)
parsed = handlers.response_parser(
response=response,
response_model=response_model,
validation_context=context,
strict=strict,
stream=stream,
is_async=True,
)
logger.debug(
f"Successfully parsed response on attempt "
f"{attempt.retry_state.attempt_number}"
)
return _finalize_parsed_response(parsed, response)
except IncompleteOutputException:
raise
except _RETRYABLE_PARSE_ERRORS as e:
logger.debug(f"Validation error on attempt {attempt_number}: {e}")
failed_attempts.append(
FailedAttempt(
attempt_number=attempt_number,
exception=e,
completion=response,
)
)
last_exception = e
if hooks:
hooks.emit_parse_error(
e,
**_attempt_metadata(
attempt_number=attempt_number,
max_attempts=max_attempts,
is_last_attempt=max_attempts == attempt_number,
),
)
# Prepare reask using registry
kwargs = handlers.reask_handler(
kwargs=kwargs,
response=response,
exception=e,
)
# Will retry with modified kwargs
raise
except IncompleteOutputException:
raise
except Exception as e:
# Max retries exceeded or non-validation error occurred
last_exception = e
logger.error(
f"Max retries exceeded. Total attempts: {last_attempt_number}, "
f"Last error: {last_exception}"
)
if hooks:
hooks.emit_completion_last_attempt(
last_exception,
**_attempt_metadata(
attempt_number=last_attempt_number or len(failed_attempts),
max_attempts=max_attempts,
is_last_attempt=True,
),
)
raise InstructorRetryException(
str(last_exception),
last_completion=failed_attempts[-1].completion if failed_attempts else None,
n_attempts=last_attempt_number,
total_usage=total_usage,
messages=extract_messages(kwargs),
create_kwargs=kwargs,
failed_attempts=failed_attempts,
) from last_exception
# Should never reach here
logger.error("Unexpected code path in retry_async_v2")
raise InstructorRetryException(
str(last_exception) if last_exception else "Unknown error",
last_completion=failed_attempts[-1].completion if failed_attempts else None,
n_attempts=last_attempt_number,
total_usage=total_usage,
messages=extract_messages(kwargs),
create_kwargs=kwargs,
failed_attempts=failed_attempts,
)