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

mirascope/llm/retries/retry_prompts.py

@@ -0,0 +1,227 @@
+"""RetryPrompt extends Prompt to return RetryResponse instead of Response."""
+
+from collections.abc import Callable
+from typing import Any, Generic, overload
+
+from ..formatting import FormatSpec, FormattableT
+from ..models import Model
+from ..prompts import (
+    AsyncPrompt,
+    Prompt,
+)
+from ..prompts.prompts import BasePrompt
+from ..prompts.protocols import AsyncMessageTemplate, MessageTemplate
+from ..providers import ModelId
+from ..tools import AsyncToolkit, Toolkit
+from ..types import P
+from .retry_config import RetryConfig
+from .retry_models import RetryModel
+from .retry_responses import (
+    AsyncRetryResponse,
+    RetryResponse,
+)
+from .retry_stream_responses import (
+    AsyncRetryStreamResponse,
+    RetryStreamResponse,
+)
+
+
+def _ensure_retry_model(model: Model | ModelId, config: RetryConfig) -> RetryModel:
+    """Ensure a model has retry capabilities, adding them if necessary."""
+    if isinstance(model, RetryModel):
+        return model
+    if isinstance(model, str):
+        model = Model(model)
+    return RetryModel(model=model, retry_config=config)
+
+
+class BaseRetryPrompt(BasePrompt[Callable[..., Any]]):
+    """Base class for retry-enabled prompts that adds retry_config."""
+
+    retry_config: RetryConfig
+    """Configuration for retry behavior."""
+
+
+class RetryPrompt(BaseRetryPrompt, Prompt[P, FormattableT], Generic[P, FormattableT]):
+    """A retry-enabled prompt that extends Prompt and returns RetryResponse.
+
+    This extends Prompt and overrides call methods to return RetryResponse
+    instead of Response. It handles wrapping the provided Model in a RetryModel
+    if necessary.
+    """
+
+    def __init__(
+        self,
+        *,
+        fn: MessageTemplate[P],
+        toolkit: Toolkit,
+        format: FormatSpec[FormattableT] | None,
+        retry_config: RetryConfig,
+    ) -> None:
+        super().__init__(fn=fn, toolkit=toolkit, format=format)
+        self.retry_config = retry_config
+
+    @overload
+    def __call__(
+        self: "RetryPrompt[P, None]",
+        model: Model | ModelId,
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> RetryResponse[None]: ...
+
+    @overload
+    def __call__(
+        self: "RetryPrompt[P, FormattableT]",
+        model: Model | ModelId,
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> RetryResponse[FormattableT]: ...
+
+    def __call__(
+        self, model: Model | ModelId, *args: P.args, **kwargs: P.kwargs
+    ) -> RetryResponse[None] | RetryResponse[FormattableT]:
+        """Generates a retry response using the provided model."""
+        return self.call(model, *args, **kwargs)
+
+    @overload
+    def call(
+        self: "RetryPrompt[P, None]",
+        model: Model | ModelId,
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> RetryResponse[None]: ...
+
+    @overload
+    def call(
+        self: "RetryPrompt[P, FormattableT]",
+        model: Model | ModelId,
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> RetryResponse[FormattableT]: ...
+
+    def call(
+        self, model: Model | ModelId, *args: P.args, **kwargs: P.kwargs
+    ) -> RetryResponse[None] | RetryResponse[FormattableT]:
+        """Generates a retry response using the provided model."""
+        retry_model = _ensure_retry_model(model, self.retry_config)
+        messages = self.messages(*args, **kwargs)
+        return retry_model.call(
+            content=messages, tools=self.toolkit, format=self.format
+        )
+
+    @overload
+    def stream(
+        self: "RetryPrompt[P, None]",
+        model: Model | ModelId,
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> RetryStreamResponse[None]: ...
+
+    @overload
+    def stream(
+        self: "RetryPrompt[P, FormattableT]",
+        model: Model | ModelId,
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> RetryStreamResponse[FormattableT]: ...
+
+    def stream(
+        self, model: Model | ModelId, *args: P.args, **kwargs: P.kwargs
+    ) -> RetryStreamResponse[None] | RetryStreamResponse[FormattableT]:
+        """Generates a retry stream response using the provided model."""
+        retry_model = _ensure_retry_model(model, self.retry_config)
+        messages = self.messages(*args, **kwargs)
+        return retry_model.stream(
+            content=messages, tools=self.toolkit, format=self.format
+        )
+
+
+class AsyncRetryPrompt(
+    BaseRetryPrompt, AsyncPrompt[P, FormattableT], Generic[P, FormattableT]
+):
+    """An async retry-enabled prompt that extends AsyncPrompt and returns AsyncRetryResponse."""
+
+    def __init__(
+        self,
+        *,
+        fn: AsyncMessageTemplate[P],
+        toolkit: AsyncToolkit,
+        format: FormatSpec[FormattableT] | None,
+        retry_config: RetryConfig,
+    ) -> None:
+        super().__init__(fn=fn, toolkit=toolkit, format=format)
+        self.retry_config = retry_config
+
+    @overload
+    async def __call__(
+        self: "AsyncRetryPrompt[P, None]",
+        model: Model | ModelId,
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> AsyncRetryResponse[None]: ...
+
+    @overload
+    async def __call__(
+        self: "AsyncRetryPrompt[P, FormattableT]",
+        model: Model | ModelId,
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> AsyncRetryResponse[FormattableT]: ...
+
+    async def __call__(
+        self, model: Model | ModelId, *args: P.args, **kwargs: P.kwargs
+    ) -> AsyncRetryResponse[None] | AsyncRetryResponse[FormattableT]:
+        """Generates a retry response using the provided model asynchronously."""
+        return await self.call(model, *args, **kwargs)
+
+    @overload
+    async def call(
+        self: "AsyncRetryPrompt[P, None]",
+        model: Model | ModelId,
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> AsyncRetryResponse[None]: ...
+
+    @overload
+    async def call(
+        self: "AsyncRetryPrompt[P, FormattableT]",
+        model: Model | ModelId,
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> AsyncRetryResponse[FormattableT]: ...
+
+    async def call(
+        self, model: Model | ModelId, *args: P.args, **kwargs: P.kwargs
+    ) -> AsyncRetryResponse[None] | AsyncRetryResponse[FormattableT]:
+        """Generates a retry response using the provided model asynchronously."""
+        retry_model = _ensure_retry_model(model, self.retry_config)
+        messages = await self.messages(*args, **kwargs)
+        return await retry_model.call_async(
+            content=messages, tools=self.toolkit, format=self.format
+        )
+
+    @overload
+    async def stream(
+        self: "AsyncRetryPrompt[P, None]",
+        model: Model | ModelId,
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> AsyncRetryStreamResponse[None]: ...
+
+    @overload
+    async def stream(
+        self: "AsyncRetryPrompt[P, FormattableT]",
+        model: Model | ModelId,
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> AsyncRetryStreamResponse[FormattableT]: ...
+
+    async def stream(
+        self, model: Model | ModelId, *args: P.args, **kwargs: P.kwargs
+    ) -> AsyncRetryStreamResponse[None] | AsyncRetryStreamResponse[FormattableT]:
+        """Generates a retry stream response using the provided model asynchronously."""
+        retry_model = _ensure_retry_model(model, self.retry_config)
+        messages = await self.messages(*args, **kwargs)
+        return await retry_model.stream_async(
+            content=messages, tools=self.toolkit, format=self.format
+        )

mirascope/llm/retries/retry_responses.py

@@ -0,0 +1,340 @@
+"""Retry response classes that extend base responses with retry metadata."""
+
+from typing import TYPE_CHECKING, Generic, overload
+
+from ..context import Context, DepsT
+from ..formatting import FormattableT
+from ..messages import UserContent
+from ..responses import (
+    AsyncContextResponse,
+    AsyncResponse,
+    ContextResponse,
+    Response,
+)
+from .utils import RetryFailure, get_retry_model_from_context
+
+if TYPE_CHECKING:
+    from .retry_config import RetryConfig
+    from .retry_models import RetryModel
+
+
+class RetryResponse(Response[FormattableT]):
+    """Response that includes retry metadata.
+
+    Extends `Response` directly, copying all attributes from a wrapped response
+    and adding retry configuration. The `model` property returns a `RetryModel`
+    with the successful model as the active model.
+    """
+
+    _retry_model: "RetryModel"
+    """The RetryModel with the successful model as active."""
+
+    retry_failures: list[RetryFailure]
+    """Failed attempts before the successful one (empty if first attempt succeeded)."""
+
+    def __init__(
+        self,
+        response: Response[FormattableT],
+        retry_model: "RetryModel",
+        retry_failures: list[RetryFailure],
+    ) -> None:
+        """Initialize a RetryResponse.
+
+        Args:
+            response: The successful response from the LLM.
+            retry_model: RetryModel with the successful model as active.
+            retry_failures: List of failed attempts before success.
+        """
+        # Copy all attributes from the wrapped response
+        for key, value in response.__dict__.items():
+            object.__setattr__(self, key, value)
+        self._retry_model = retry_model
+        self.retry_failures = retry_failures
+
+    @property
+    def model(self) -> "RetryModel":
+        """A RetryModel with the successful model as active.
+
+        If a model is set in context (via `llm.model()` or `llm.retry_model()`),
+        that model is used instead, wrapped with this response's retry config.
+
+        Otherwise, this RetryModel has the same pool of available models, but with
+        the model that succeeded set as the active model. This means:
+        - `response.model.model_id` equals `response.model_id`
+        - `response.resume()` will try the successful model first
+        """
+        return get_retry_model_from_context(self._retry_model)
+
+    @property
+    def retry_config(self) -> "RetryConfig":
+        """The retry configuration for this response."""
+        return self.model.retry_config
+
+    @overload
+    def resume(
+        self: "RetryResponse[None]", content: UserContent
+    ) -> "RetryResponse[None]": ...
+
+    @overload
+    def resume(
+        self: "RetryResponse[FormattableT]", content: UserContent
+    ) -> "RetryResponse[FormattableT]": ...
+
+    def resume(
+        self, content: UserContent
+    ) -> "RetryResponse[None] | RetryResponse[FormattableT]":
+        """Generate a new RetryResponse using this response's messages with additional user content.
+
+        Args:
+            content: The new user message content to append to the message history.
+
+        Returns:
+            A new RetryResponse instance generated from the extended message history.
+        """
+        return self.model.resume(response=self, content=content)
+
+
+class AsyncRetryResponse(AsyncResponse[FormattableT]):
+    """Async response that includes retry metadata.
+
+    Extends `AsyncResponse` directly, copying all attributes from a wrapped response
+    and adding retry configuration. The `model` property returns a `RetryModel`
+    with the successful model as the active model.
+    """
+
+    _retry_model: "RetryModel"
+    """The RetryModel with the successful model as active."""
+
+    retry_failures: list[RetryFailure]
+    """Failed attempts before the successful one (empty if first attempt succeeded)."""
+
+    def __init__(
+        self,
+        response: AsyncResponse[FormattableT],
+        retry_model: "RetryModel",
+        retry_failures: list[RetryFailure],
+    ) -> None:
+        """Initialize an AsyncRetryResponse.
+
+        Args:
+            response: The successful async response from the LLM.
+            retry_model: RetryModel with the successful model as active.
+            retry_failures: List of failed attempts before success.
+        """
+        # Copy all attributes from the wrapped response
+        for key, value in response.__dict__.items():
+            object.__setattr__(self, key, value)
+        self._retry_model = retry_model
+        self.retry_failures = retry_failures
+
+    @property
+    def model(self) -> "RetryModel":
+        """A RetryModel with the successful model as active.
+
+        If a model is set in context (via `llm.model()` or `llm.retry_model()`),
+        that model is used instead, wrapped with this response's retry config.
+
+        Otherwise, this RetryModel has the same pool of available models, but with
+        the model that succeeded set as the active model. This means:
+        - `response.model.model_id` equals `response.model_id`
+        - `response.resume()` will try the successful model first
+        """
+        return get_retry_model_from_context(self._retry_model)
+
+    @property
+    def retry_config(self) -> "RetryConfig":
+        """The retry configuration for this response."""
+        return self.model.retry_config
+
+    @overload
+    async def resume(
+        self: "AsyncRetryResponse[None]", content: UserContent
+    ) -> "AsyncRetryResponse[None]": ...
+
+    @overload
+    async def resume(
+        self: "AsyncRetryResponse[FormattableT]", content: UserContent
+    ) -> "AsyncRetryResponse[FormattableT]": ...
+
+    async def resume(
+        self, content: UserContent
+    ) -> "AsyncRetryResponse[None] | AsyncRetryResponse[FormattableT]":
+        """Generate a new AsyncRetryResponse using this response's messages with additional user content.
+
+        Args:
+            content: The new user message content to append to the message history.
+
+        Returns:
+            A new AsyncRetryResponse instance generated from the extended message history.
+        """
+        return await self.model.resume_async(response=self, content=content)
+
+
+class ContextRetryResponse(
+    ContextResponse[DepsT, FormattableT], Generic[DepsT, FormattableT]
+):
+    """Context response that includes retry metadata.
+
+    Extends `ContextResponse` directly, copying all attributes from a wrapped response
+    and adding retry configuration. The `model` property returns a `RetryModel`
+    with the successful model as the active model.
+    """
+
+    _retry_model: "RetryModel"
+    """The RetryModel with the successful model as active."""
+
+    retry_failures: list[RetryFailure]
+    """Failed attempts before the successful one (empty if first attempt succeeded)."""
+
+    def __init__(
+        self,
+        response: ContextResponse[DepsT, FormattableT],
+        retry_model: "RetryModel",
+        retry_failures: list[RetryFailure],
+    ) -> None:
+        """Initialize a ContextRetryResponse.
+
+        Args:
+            response: The successful context response from the LLM.
+            retry_model: RetryModel with the successful model as active.
+            retry_failures: List of failed attempts before success.
+        """
+        # Copy all attributes from the wrapped response
+        for key, value in response.__dict__.items():
+            object.__setattr__(self, key, value)
+        self._retry_model = retry_model
+        self.retry_failures = retry_failures
+
+    @property
+    def model(self) -> "RetryModel":
+        """A RetryModel with the successful model as active.
+
+        If a model is set in context (via `llm.model()` or `llm.retry_model()`),
+        that model is used instead, wrapped with this response's retry config.
+
+        Otherwise, this RetryModel has the same pool of available models, but with
+        the model that succeeded set as the active model. This means:
+        - `response.model.model_id` equals `response.model_id`
+        - `response.resume()` will try the successful model first
+        """
+        return get_retry_model_from_context(self._retry_model)
+
+    @property
+    def retry_config(self) -> "RetryConfig":
+        """The retry configuration for this response."""
+        return self.model.retry_config
+
+    @overload
+    def resume(
+        self: "ContextRetryResponse[DepsT, None]",
+        ctx: Context[DepsT],
+        content: UserContent,
+    ) -> "ContextRetryResponse[DepsT, None]": ...
+
+    @overload
+    def resume(
+        self: "ContextRetryResponse[DepsT, FormattableT]",
+        ctx: Context[DepsT],
+        content: UserContent,
+    ) -> "ContextRetryResponse[DepsT, FormattableT]": ...
+
+    def resume(
+        self, ctx: Context[DepsT], content: UserContent
+    ) -> (
+        "ContextRetryResponse[DepsT, None] | ContextRetryResponse[DepsT, FormattableT]"
+    ):
+        """Generate a new ContextRetryResponse using this response's messages with additional user content.
+
+        Args:
+            ctx: A `Context` with the required deps type.
+            content: The new user message content to append to the message history.
+
+        Returns:
+            A new ContextRetryResponse instance generated from the extended message history.
+        """
+        return self.model.context_resume(ctx=ctx, response=self, content=content)
+
+
+class AsyncContextRetryResponse(
+    AsyncContextResponse[DepsT, FormattableT], Generic[DepsT, FormattableT]
+):
+    """Async context response that includes retry metadata.
+
+    Extends `AsyncContextResponse` directly, copying all attributes from a wrapped response
+    and adding retry configuration. The `model` property returns a `RetryModel`
+    with the successful model as the active model.
+    """
+
+    _retry_model: "RetryModel"
+    """The RetryModel with the successful model as active."""
+
+    retry_failures: list[RetryFailure]
+    """Failed attempts before the successful one (empty if first attempt succeeded)."""
+
+    def __init__(
+        self,
+        response: AsyncContextResponse[DepsT, FormattableT],
+        retry_model: "RetryModel",
+        retry_failures: list[RetryFailure],
+    ) -> None:
+        """Initialize an AsyncContextRetryResponse.
+
+        Args:
+            response: The successful async context response from the LLM.
+            retry_model: RetryModel with the successful model as active.
+            retry_failures: List of failed attempts before success.
+        """
+        # Copy all attributes from the wrapped response
+        for key, value in response.__dict__.items():
+            object.__setattr__(self, key, value)
+        self._retry_model = retry_model
+        self.retry_failures = retry_failures
+
+    @property
+    def model(self) -> "RetryModel":
+        """A RetryModel with the successful model as active.
+
+        If a model is set in context (via `llm.model()` or `llm.retry_model()`),
+        that model is used instead, wrapped with this response's retry config.
+
+        Otherwise, this RetryModel has the same pool of available models, but with
+        the model that succeeded set as the active model. This means:
+        - `response.model.model_id` equals `response.model_id`
+        - `response.resume()` will try the successful model first
+        """
+        return get_retry_model_from_context(self._retry_model)
+
+    @property
+    def retry_config(self) -> "RetryConfig":
+        """The retry configuration for this response."""
+        return self.model.retry_config
+
+    @overload
+    async def resume(
+        self: "AsyncContextRetryResponse[DepsT, None]",
+        ctx: Context[DepsT],
+        content: UserContent,
+    ) -> "AsyncContextRetryResponse[DepsT, None]": ...
+
+    @overload
+    async def resume(
+        self: "AsyncContextRetryResponse[DepsT, FormattableT]",
+        ctx: Context[DepsT],
+        content: UserContent,
+    ) -> "AsyncContextRetryResponse[DepsT, FormattableT]": ...
+
+    async def resume(
+        self, ctx: Context[DepsT], content: UserContent
+    ) -> "AsyncContextRetryResponse[DepsT, None] | AsyncContextRetryResponse[DepsT, FormattableT]":
+        """Generate a new AsyncContextRetryResponse using this response's messages with additional user content.
+
+        Args:
+            ctx: A `Context` with the required deps type.
+            content: The new user message content to append to the message history.
+
+        Returns:
+            A new AsyncContextRetryResponse instance generated from the extended message history.
+        """
+        return await self.model.context_resume_async(
+            ctx=ctx, response=self, content=content
+        )