mirascope 2.1.0__py3-none-any.whl → 2.2.0__py3-none-any.whl

mirascope/llm/retries/retry_calls.py

@@ -0,0 +1,159 @@
+"""RetryCall extends Call to return RetryResponse instead of Response."""
+
+from typing import Generic, TypeVar, overload
+
+from ..calls.calls import BaseCall
+from ..formatting import FormattableT
+from ..types import P
+from .retry_config import RetryConfig
+from .retry_models import RetryModel
+from .retry_prompts import (
+    AsyncRetryPrompt,
+    BaseRetryPrompt,
+    RetryPrompt,
+)
+from .retry_responses import (
+    AsyncRetryResponse,
+    RetryResponse,
+)
+from .retry_stream_responses import (
+    AsyncRetryStreamResponse,
+    RetryStreamResponse,
+)
+
+RetryPromptT = TypeVar("RetryPromptT", bound=BaseRetryPrompt)
+
+
+class BaseRetryCall(BaseCall[RetryPromptT], Generic[RetryPromptT]):
+    """Base class for retry-enabled calls with shared model wrapping."""
+
+    @property
+    def retry_config(self) -> RetryConfig:
+        """The retry configuration from the prompt."""
+        return self.prompt.retry_config
+
+    @property
+    def model(self) -> RetryModel:
+        """The model used for generating responses. May be overwritten via `with llm.model(...)`."""
+        model = super().model
+        if isinstance(model, RetryModel):
+            return model
+        return RetryModel(model, self.retry_config)
+
+
+class RetryCall(BaseRetryCall[RetryPrompt[P, FormattableT]], Generic[P, FormattableT]):
+    """A retry-enabled call that extends BaseRetryCall and returns RetryResponse.
+
+    This extends BaseRetryCall with a RetryPrompt and returns RetryResponse
+    instead of Response. The prompt contains the retry configuration.
+
+    The model can be overridden at runtime using `with llm.model(...)` context manager.
+    """
+
+    @overload
+    def __call__(
+        self: "RetryCall[P, None]", *args: P.args, **kwargs: P.kwargs
+    ) -> RetryResponse[None]: ...
+
+    @overload
+    def __call__(
+        self: "RetryCall[P, FormattableT]", *args: P.args, **kwargs: P.kwargs
+    ) -> RetryResponse[FormattableT]: ...
+
+    def __call__(
+        self, *args: P.args, **kwargs: P.kwargs
+    ) -> RetryResponse[None] | RetryResponse[FormattableT]:
+        """Generates a retry response using the LLM."""
+        return self.call(*args, **kwargs)
+
+    @overload
+    def call(
+        self: "RetryCall[P, None]", *args: P.args, **kwargs: P.kwargs
+    ) -> RetryResponse[None]: ...
+
+    @overload
+    def call(
+        self: "RetryCall[P, FormattableT]", *args: P.args, **kwargs: P.kwargs
+    ) -> RetryResponse[FormattableT]: ...
+
+    def call(
+        self, *args: P.args, **kwargs: P.kwargs
+    ) -> RetryResponse[None] | RetryResponse[FormattableT]:
+        """Generates a retry response using the LLM."""
+        return self.prompt.call(self.model, *args, **kwargs)
+
+    @overload
+    def stream(
+        self: "RetryCall[P, None]", *args: P.args, **kwargs: P.kwargs
+    ) -> RetryStreamResponse[None]: ...
+
+    @overload
+    def stream(
+        self: "RetryCall[P, FormattableT]", *args: P.args, **kwargs: P.kwargs
+    ) -> RetryStreamResponse[FormattableT]: ...
+
+    def stream(
+        self, *args: P.args, **kwargs: P.kwargs
+    ) -> RetryStreamResponse[None] | RetryStreamResponse[FormattableT]:
+        """Generates a retry stream response using the LLM."""
+        return self.prompt.stream(self.model, *args, **kwargs)
+
+
+class AsyncRetryCall(
+    BaseRetryCall[AsyncRetryPrompt[P, FormattableT]], Generic[P, FormattableT]
+):
+    """An async retry-enabled call that extends BaseRetryCall and returns AsyncRetryResponse.
+
+    This extends BaseRetryCall with an AsyncRetryPrompt and returns AsyncRetryResponse
+    instead of AsyncResponse. The prompt contains the retry configuration.
+
+    The model can be overridden at runtime using `with llm.model(...)` context manager.
+    """
+
+    @overload
+    async def __call__(
+        self: "AsyncRetryCall[P, None]", *args: P.args, **kwargs: P.kwargs
+    ) -> AsyncRetryResponse[None]: ...
+
+    @overload
+    async def __call__(
+        self: "AsyncRetryCall[P, FormattableT]", *args: P.args, **kwargs: P.kwargs
+    ) -> AsyncRetryResponse[FormattableT]: ...
+
+    async def __call__(
+        self, *args: P.args, **kwargs: P.kwargs
+    ) -> AsyncRetryResponse[None] | AsyncRetryResponse[FormattableT]:
+        """Generates a retry response using the LLM asynchronously."""
+        return await self.call(*args, **kwargs)
+
+    @overload
+    async def call(
+        self: "AsyncRetryCall[P, None]", *args: P.args, **kwargs: P.kwargs
+    ) -> AsyncRetryResponse[None]: ...
+
+    @overload
+    async def call(
+        self: "AsyncRetryCall[P, FormattableT]", *args: P.args, **kwargs: P.kwargs
+    ) -> AsyncRetryResponse[FormattableT]: ...
+
+    async def call(
+        self, *args: P.args, **kwargs: P.kwargs
+    ) -> AsyncRetryResponse[None] | AsyncRetryResponse[FormattableT]:
+        """Generates a retry response using the LLM asynchronously."""
+        return await self.prompt.call(self.model, *args, **kwargs)
+
+    @overload
+    async def stream(
+        self: "AsyncRetryCall[P, None]", *args: P.args, **kwargs: P.kwargs
+    ) -> AsyncRetryStreamResponse[None]: ...
+
+    @overload
+    async def stream(
+        self: "AsyncRetryCall[P, FormattableT]", *args: P.args, **kwargs: P.kwargs
+    ) -> AsyncRetryStreamResponse[FormattableT]: ...
+
+    async def stream(
+        self, *args: P.args, **kwargs: P.kwargs
+    ) -> AsyncRetryStreamResponse[None] | AsyncRetryStreamResponse[FormattableT]:
+        """Generates a retry stream response using the LLM asynchronously."""
+        return await self.prompt.stream(self.model, *args, **kwargs)

mirascope/llm/retries/retry_config.py

@@ -0,0 +1,168 @@
+"""Configuration for retry behavior."""
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, TypedDict
+from typing_extensions import Unpack
+
+from ..exceptions import (
+    ConnectionError,
+    RateLimitError,
+    ServerError,
+    TimeoutError,
+)
+
+if TYPE_CHECKING:
+    from ..models import Model
+    from ..providers import ModelId
+
+DEFAULT_RETRYABLE_ERRORS: tuple[type[Exception], ...] = (
+    ConnectionError,
+    RateLimitError,
+    ServerError,
+    TimeoutError,
+)
+"""Default exceptions that trigger a retry.
+
+These are transient errors that are likely to succeed on retry:
+- ConnectionError: Network issues, DNS failures
+- RateLimitError: Rate limits exceeded (429)
+- ServerError: Provider-side errors (500+)
+- TimeoutError: Request timeouts
+"""
+
+DEFAULT_MAX_RETRIES: int = 3
+"""Default maximum number of retries after the initial attempt fails."""
+
+DEFAULT_INITIAL_DELAY: float = 0.5
+"""Default initial delay in seconds before the first retry."""
+
+DEFAULT_MAX_DELAY: float = 60.0
+"""Default maximum delay in seconds between retries."""
+
+DEFAULT_BACKOFF_MULTIPLIER: float = 2.0
+"""Default multiplier for exponential backoff (delay *= multiplier after each retry)."""
+
+DEFAULT_JITTER: float = 0.0
+"""Default jitter factor (0.0 to 1.0) to add randomness to delays.
+
+A jitter of 0.1 means +/- 10% random variation on the calculated delay.
+"""
+
+
+class RetryArgs(TypedDict, total=False):
+    """Arguments for configuring retry behavior.
+
+    This TypedDict is used for the user-facing API where all fields are optional.
+    Use RetryConfig for the internal representation with defaults applied.
+
+    Attributes:
+        max_retries: Maximum number of retries after the initial attempt fails.
+            Defaults to 3.
+        retry_on: Tuple of exception types that should trigger a retry.
+            Defaults to DEFAULT_RETRYABLE_ERRORS (ConnectionError, RateLimitError,
+            ServerError, TimeoutError).
+        initial_delay: Initial delay in seconds before the first retry. Defaults to 0.5.
+        max_delay: Maximum delay in seconds between retries. Defaults to 60.0.
+        backoff_multiplier: Multiplier for exponential backoff. Defaults to 2.0.
+        jitter: Jitter factor (0.0 to 1.0) to add randomness to delays. Defaults to 0.0.
+    """
+
+    max_retries: int
+    """Maximum number of retries after the initial attempt fails. Defaults to 3."""
+
+    retry_on: tuple[type[Exception], ...]
+    """Tuple of exception types that should trigger a retry.
+
+    Defaults to DEFAULT_RETRYABLE_ERRORS (ConnectionError, RateLimitError,
+    ServerError, TimeoutError).
+    """
+
+    initial_delay: float
+    """Initial delay in seconds before the first retry. Defaults to 1.0."""
+
+    max_delay: float
+    """Maximum delay in seconds between retries. Defaults to 60.0."""
+
+    backoff_multiplier: float
+    """Multiplier for exponential backoff (delay *= multiplier after each retry). Defaults to 2.0."""
+
+    jitter: float
+    """Jitter factor (0.0 to 1.0) to add randomness to delays. Defaults to 0.0."""
+
+    fallback_models: "Sequence[Model | ModelId]"
+    """Sequence of fallback models to try if the primary model fails.
+
+    Each model gets its own full retry budget (max_retries applies per model).
+    ModelId strings inherit params from the primary model; Model instances
+    use their own params. Defaults to empty tuple (no fallbacks).
+    """
+
+
+@dataclass(frozen=True)
+class RetryConfig:
+    """Configuration for retry behavior with defaults applied.
+
+    Attributes:
+        max_retries: Maximum number of retries after the initial attempt fails.
+        retry_on: Tuple of exception types that should trigger a retry.
+        initial_delay: Initial delay in seconds before the first retry.
+        max_delay: Maximum delay in seconds between retries.
+        backoff_multiplier: Multiplier for exponential backoff.
+        jitter: Jitter factor (0.0 to 1.0) to add randomness to delays.
+    """
+
+    max_retries: int = DEFAULT_MAX_RETRIES
+    """Maximum number of retries after the initial attempt fails."""
+
+    retry_on: tuple[type[Exception], ...] = DEFAULT_RETRYABLE_ERRORS
+    """Tuple of exception types that should trigger a retry."""
+
+    fallback_models: "tuple[Model | ModelId, ...]" = ()
+    """Sequence of fallback models to try if the primary model fails.
+
+    Each model gets its own full retry budget (max_retries applies per model).
+    ModelId strings inherit params from the primary model; Model instances
+    use their own params. Defaults to empty tuple (no fallbacks).
+    """
+
+    initial_delay: float = DEFAULT_INITIAL_DELAY
+    """Initial delay in seconds before the first retry."""
+
+    max_delay: float = DEFAULT_MAX_DELAY
+    """Maximum delay in seconds between retries."""
+
+    backoff_multiplier: float = DEFAULT_BACKOFF_MULTIPLIER
+    """Multiplier for exponential backoff (delay *= multiplier after each retry)."""
+
+    jitter: float = DEFAULT_JITTER
+    """Jitter factor (0.0 to 1.0) to add randomness to delays."""
+
+    def __post_init__(self) -> None:
+        """Validate configuration values."""
+        if self.max_retries < 0:
+            raise ValueError("max_retries must be non-negative")
+        if self.initial_delay < 0:
+            raise ValueError("initial_delay must be non-negative")
+        if self.max_delay < 0:
+            raise ValueError("max_delay must be non-negative")
+        if self.backoff_multiplier < 1:
+            raise ValueError("backoff_multiplier must be >= 1")
+        if not 0 <= self.jitter <= 1:
+            raise ValueError("jitter must be between 0.0 and 1.0")
+
+    @classmethod
+    def from_args(cls, **args: Unpack[RetryArgs]) -> "RetryConfig":
+        """Create a RetryConfig from RetryArgs kwargs."""
+        fallback_models = args.get("fallback_models", ())
+        return cls(
+            max_retries=args.get("max_retries", DEFAULT_MAX_RETRIES),
+            retry_on=args.get("retry_on", DEFAULT_RETRYABLE_ERRORS),
+            fallback_models=tuple(fallback_models),
+            initial_delay=args.get("initial_delay", DEFAULT_INITIAL_DELAY),
+            max_delay=args.get("max_delay", DEFAULT_MAX_DELAY),
+            backoff_multiplier=args.get(
+                "backoff_multiplier", DEFAULT_BACKOFF_MULTIPLIER
+            ),
+            jitter=args.get("jitter", DEFAULT_JITTER),
+        )

mirascope/llm/retries/retry_decorator.py

@@ -0,0 +1,258 @@
+"""The @llm.retry decorator for adding retry logic to prompts and calls."""
+
+from typing import overload
+from typing_extensions import Unpack
+
+from ..calls import AsyncCall, AsyncContextCall, Call, ContextCall
+from ..context import DepsT
+from ..formatting import FormattableT
+from ..prompts import AsyncContextPrompt, AsyncPrompt, ContextPrompt, Prompt
+from ..types import P
+from .retry_calls import (
+    AsyncRetryCall,
+    RetryCall,
+)
+from .retry_config import RetryArgs, RetryConfig
+from .retry_prompts import (
+    AsyncRetryPrompt,
+    RetryPrompt,
+)
+
+RetryTarget = (
+    ContextPrompt[P, DepsT, FormattableT]
+    | AsyncContextPrompt[P, DepsT, FormattableT]
+    | Prompt[P, FormattableT]
+    | RetryPrompt[P, FormattableT]
+    | AsyncPrompt[P, FormattableT]
+    | AsyncRetryPrompt[P, FormattableT]
+    | ContextCall[P, DepsT, FormattableT]
+    | AsyncContextCall[P, DepsT, FormattableT]
+    | Call[P, FormattableT]
+    | RetryCall[P, FormattableT]
+    | AsyncCall[P, FormattableT]
+    | AsyncRetryCall[P, FormattableT]
+)
+"""Union type for all targets that can be wrapped with retry logic.
+
+Includes Prompts and Calls (both sync and async variants).
+For models, use `llm.retry_model()` instead.
+"""
+
+RetryResult = (
+    RetryPrompt[P, FormattableT]
+    | AsyncRetryPrompt[P, FormattableT]
+    | RetryCall[P, FormattableT]
+    | AsyncRetryCall[P, FormattableT]
+)
+"""Union type for all retry-wrapped results.
+
+When a target is wrapped with retry logic, it returns one of these types.
+"""
+
+
+# Overloads for direct wrapping: retry(target, ...)
+
+
+@overload
+def retry(
+    target: Prompt[P, FormattableT] | RetryPrompt[P, FormattableT],
+    **config: Unpack[RetryArgs],
+) -> RetryPrompt[P, FormattableT]:
+    """Wrap a Prompt with retry logic."""
+    ...
+
+
+@overload
+def retry(
+    target: AsyncPrompt[P, FormattableT] | AsyncRetryPrompt[P, FormattableT],
+    **config: Unpack[RetryArgs],
+) -> AsyncRetryPrompt[P, FormattableT]:
+    """Wrap an AsyncPrompt with retry logic."""
+    ...
+
+
+@overload
+def retry(
+    target: Call[P, FormattableT] | RetryCall[P, FormattableT],
+    **config: Unpack[RetryArgs],
+) -> RetryCall[P, FormattableT]:
+    """Wrap a Call with retry logic."""
+    ...
+
+
+@overload
+def retry(
+    target: AsyncCall[P, FormattableT] | AsyncRetryCall[P, FormattableT],
+    **config: Unpack[RetryArgs],
+) -> AsyncRetryCall[P, FormattableT]:
+    """Wrap an AsyncCall with retry logic."""
+    ...
+
+
+# Overloads for decorator usage: @retry(...)
+
+
+@overload
+def retry(
+    target: None = None,
+    **config: Unpack[RetryArgs],
+) -> "RetryDecorator":
+    """Return a decorator that adds retry logic."""
+    ...
+
+
+class RetryDecorator:
+    """A decorator class that wraps targets with retry logic.
+
+    This class is returned when `retry()` is called without a target,
+    allowing it to be used as a parameterized decorator.
+    """
+
+    def __init__(self, config: RetryConfig) -> None:
+        self._config = config
+
+    @overload
+    def __call__(
+        self, target: Prompt[P, FormattableT] | RetryPrompt[P, FormattableT]
+    ) -> RetryPrompt[P, FormattableT]: ...
+
+    @overload
+    def __call__(
+        self, target: AsyncPrompt[P, FormattableT] | AsyncRetryPrompt[P, FormattableT]
+    ) -> AsyncRetryPrompt[P, FormattableT]: ...
+
+    @overload
+    def __call__(
+        self, target: Call[P, FormattableT] | RetryCall[P, FormattableT]
+    ) -> RetryCall[P, FormattableT]: ...
+
+    @overload
+    def __call__(
+        self, target: AsyncCall[P, FormattableT] | AsyncRetryCall[P, FormattableT]
+    ) -> AsyncRetryCall[P, FormattableT]: ...
+
+    def __call__(
+        self,
+        target: RetryTarget[P, DepsT, FormattableT],
+    ) -> RetryResult[P, FormattableT]:
+        """Apply retry logic to the target."""
+        return _wrap_target(target, self._config)
+
+
+def _wrap_target(
+    target: RetryTarget[P, DepsT, FormattableT],
+    retry_config: RetryConfig,
+) -> RetryResult[P, FormattableT]:
+    """Internal function to wrap a target with retry logic."""
+    # Wrap RetryCall variants first - they inherit from BaseCall, not Call/AsyncCall
+    if isinstance(target, AsyncRetryCall):
+        prompt = target.prompt
+        return AsyncRetryCall(
+            default_model=target.default_model,
+            prompt=AsyncRetryPrompt(
+                fn=prompt.fn,
+                toolkit=prompt.toolkit,
+                format=prompt.format,
+                retry_config=retry_config,
+            ),
+        )
+    if isinstance(target, RetryCall):
+        prompt = target.prompt
+        return RetryCall(
+            default_model=target.default_model,
+            prompt=RetryPrompt(
+                fn=prompt.fn,
+                toolkit=prompt.toolkit,
+                format=prompt.format,
+                retry_config=retry_config,
+            ),
+        )
+
+    # Wrap Call variants - AsyncCall check must come before Call
+    if isinstance(target, AsyncCall):
+        prompt = target.prompt
+        return AsyncRetryCall(
+            default_model=target.default_model,
+            prompt=AsyncRetryPrompt(
+                fn=prompt.fn,
+                toolkit=prompt.toolkit,
+                format=prompt.format,
+                retry_config=retry_config,
+            ),
+        )
+    if isinstance(target, Call):
+        prompt = target.prompt
+        return RetryCall(
+            default_model=target.default_model,
+            prompt=RetryPrompt(
+                fn=prompt.fn,
+                toolkit=prompt.toolkit,
+                format=prompt.format,
+                retry_config=retry_config,
+            ),
+        )
+
+    # Wrap Prompt variants - AsyncPrompt check must come before Prompt
+    if isinstance(target, AsyncPrompt):
+        return AsyncRetryPrompt(
+            fn=target.fn,
+            toolkit=target.toolkit,
+            format=target.format,
+            retry_config=retry_config,
+        )
+    if isinstance(target, Prompt):
+        return RetryPrompt(
+            fn=target.fn,
+            toolkit=target.toolkit,
+            format=target.format,
+            retry_config=retry_config,
+        )
+
+    raise ValueError(f"Unsupported target type for retry: {type(target)}")
+
+
+def retry(
+    target: RetryTarget[P, DepsT, FormattableT] | None = None,
+    **config: Unpack[RetryArgs],
+) -> RetryResult[P, FormattableT] | RetryDecorator:
+    """Add retry logic to a Prompt or Call.
+
+    This function can be used in two ways:
+
+    1. Direct wrapping:
+        ```python
+        retry_call = llm.retry(my_call, max_retries=2)
+        ```
+
+    2. As a decorator:
+        ```python
+        @llm.retry(max_retries=2)
+        @llm.call("openai/gpt-4")
+        def my_call() -> str:
+            return "Hello"
+        ```
+
+    This function wraps the provided target with retry capabilities, allowing it to
+    automatically retry on failures, handle rate limits, and use fallback models.
+    The retry configuration can be provided as keyword arguments.
+
+    For models, use `llm.retry_model()` instead.
+
+    If a RetryPrompt or RetryCall is passed, then this will return a new
+    RetryPrompt or RetryCall using the retry_config passed to the function.
+
+    Args:
+        target: The Prompt or Call to wrap with retry logic. If None,
+            returns a decorator that can be applied to a target.
+        **config: Configuration for retry behavior (see RetryArgs).
+
+    Returns:
+        A retry-wrapped version of the target that returns Retry* response types,
+        or a RetryDecorator if no target is provided.
+    """
+    retry_config = RetryConfig.from_args(**config)
+
+    if target is None:
+        return RetryDecorator(retry_config)
+
+    return _wrap_target(target, retry_config)