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

mirascope/llm/retries/utils.py

@@ -0,0 +1,159 @@
+"""Utility functions for retry logic."""
+
+import random
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, TypeVar
+
+from ..exceptions import RetriesExhausted
+from .retry_config import RetryConfig
+
+if TYPE_CHECKING:
+    from ..models import Model
+    from .retry_models import RetryModel
+
+
+def get_retry_model_from_context(stored_retry_model: "RetryModel") -> "RetryModel":
+    """Get the RetryModel to use, checking context for overrides.
+
+    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.
+
+    Args:
+        stored_retry_model: The RetryModel stored on the response.
+
+    Returns:
+        Either the context model (wrapped in RetryModel if needed) or the stored model.
+    """
+    from ..models import model_from_context
+    from .retry_models import RetryModel
+
+    context_model = model_from_context()
+    if context_model is not None:
+        if isinstance(context_model, RetryModel):
+            return context_model
+        return RetryModel(context_model, stored_retry_model.retry_config)
+    return stored_retry_model
+
+
+_ResultT = TypeVar("_ResultT")
+
+
+@dataclass
+class RetryFailure:
+    """A failed attempt to call a model.
+
+    Attributes:
+        model: The model that was tried.
+        exception: The exception that was raised.
+    """
+
+    model: "Model"
+    exception: BaseException
+
+
+def calculate_delay(config: RetryConfig, attempt_for_model: int) -> float:
+    """Calculate the delay before the next retry attempt.
+
+    Args:
+        config: The retry configuration with backoff settings.
+        attempt_for_model: The retry attempt number for this model (1-indexed).
+
+    Returns:
+        The delay in seconds, with exponential backoff, capped at max_delay,
+        and optionally with jitter applied.
+    """
+    # Calculate base delay with exponential backoff
+    delay = config.initial_delay * (
+        config.backoff_multiplier ** (attempt_for_model - 1)
+    )
+
+    # Cap at max_delay
+    delay = min(delay, config.max_delay)
+
+    # Apply jitter if configured
+    if config.jitter > 0:
+        jitter_range = delay * config.jitter
+        delay = delay + random.uniform(-jitter_range, jitter_range)
+        # Ensure delay doesn't go negative
+        delay = max(0, delay)
+
+    return delay
+
+
+def with_retry(
+    fn: Callable[["Model"], _ResultT],
+    retry_model: "RetryModel",
+) -> tuple[_ResultT, list[RetryFailure], "RetryModel"]:
+    """Execute a function with retry logic across the RetryModel's models.
+
+    Tries the active model first, then fallbacks. Each model gets its own
+    full retry budget. Returns an updated RetryModel with the successful
+    model set as active.
+
+    Args:
+        fn: Function that takes a Model and returns a result.
+        retry_model: The RetryModel containing models and retry config.
+
+    Returns:
+        A tuple of (result, failures, updated_retry_model).
+        failures contains all failed attempts before success.
+        The updated_retry_model has the successful model as active.
+
+    Raises:
+        Exception: The last exception encountered if all models exhaust retries.
+    """
+    config = retry_model.retry_config
+    failures: list[RetryFailure] = []
+
+    for variant in retry_model.variants():
+        model = variant.get_active_model()
+        try:
+            result = fn(model)
+            return result, failures, variant
+        except config.retry_on as e:
+            failures.append(RetryFailure(model=model, exception=e))
+
+    # All models exhausted
+    if failures:
+        raise RetriesExhausted(failures)
+    raise AssertionError("Unreachable: no models provided")  # pragma: no cover
+
+
+async def with_retry_async(
+    fn: Callable[["Model"], Awaitable[_ResultT]],
+    retry_model: "RetryModel",
+) -> tuple[_ResultT, list[RetryFailure], "RetryModel"]:
+    """Execute an async function with retry logic across the RetryModel's models.
+
+    Tries the active model first, then fallbacks. Each model gets its own
+    full retry budget. Returns an updated RetryModel with the successful
+    model set as active.
+
+    Args:
+        fn: Async function that takes a Model and returns a result.
+        retry_model: The RetryModel containing models and retry config.
+
+    Returns:
+        A tuple of (result, failures, updated_retry_model).
+        failures contains all failed attempts before success.
+        The updated_retry_model has the successful model as active.
+
+    Raises:
+        Exception: The last exception encountered if all models exhaust retries.
+    """
+    config = retry_model.retry_config
+    failures: list[RetryFailure] = []
+
+    async for variant in retry_model.variants_async():
+        model = variant.get_active_model()
+        try:
+            result = await fn(model)
+            return result, failures, variant
+        except config.retry_on as e:
+            failures.append(RetryFailure(model=model, exception=e))
+
+    # All models exhausted
+    if failures:
+        raise RetriesExhausted(failures)
+    raise AssertionError("Unreachable: no models provided")  # pragma: no cover

mirascope/ops/_internal/versioned_calls.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 
 import logging
 from dataclasses import dataclass, field
-from typing import Concatenate, Generic
+from typing import Concatenate, Generic, cast, overload
 from typing_extensions import TypeIs
 
 from ...llm.calls import AsyncCall, AsyncContextCall, Call, ContextCall
@@ -238,23 +238,61 @@ class VersionedCall(_BaseVersionedCall, Generic[P, FormattableT]):
         )
         self._compute_closure(self._call, self.call, self.stream)
 
+    @overload
+    def __call__(
+        self: VersionedCall[P, None], *args: P.args, **kwargs: P.kwargs
+    ) -> Response: ...
+
+    @overload
+    def __call__(
+        self: VersionedCall[P, FormattableT], *args: P.args, **kwargs: P.kwargs
+    ) -> Response[FormattableT]: ...
+
     def __call__(
         self, *args: P.args, **kwargs: P.kwargs
     ) -> Response | Response[FormattableT]:
         """Call the versioned function and return Response directly."""
         return self.call(*args, **kwargs)
 
+    @overload
+    def wrapped(
+        self: VersionedCall[P, None], *args: P.args, **kwargs: P.kwargs
+    ) -> VersionedResult[Response]: ...
+
+    @overload
+    def wrapped(
+        self: VersionedCall[P, FormattableT], *args: P.args, **kwargs: P.kwargs
+    ) -> VersionedResult[Response[FormattableT]]: ...
+
     def wrapped(
         self, *args: P.args, **kwargs: P.kwargs
-    ) -> VersionedResult[Response | Response[FormattableT]]:
+    ) -> VersionedResult[Response] | VersionedResult[Response[FormattableT]]:
         """Call the versioned function and return a wrapped Response."""
-        return self.call.wrapped(*args, **kwargs)
+        return cast(
+            "VersionedResult[Response] | VersionedResult[Response[FormattableT]]",
+            self.call.wrapped(*args, **kwargs),
+        )
+
+    @overload
+    def wrapped_stream(
+        self: VersionedCall[P, None], *args: P.args, **kwargs: P.kwargs
+    ) -> VersionedResult[StreamResponse]: ...
+
+    @overload
+    def wrapped_stream(
+        self: VersionedCall[P, FormattableT], *args: P.args, **kwargs: P.kwargs
+    ) -> VersionedResult[StreamResponse[FormattableT]]: ...
 
     def wrapped_stream(
         self, *args: P.args, **kwargs: P.kwargs
-    ) -> VersionedResult[StreamResponse | StreamResponse[FormattableT]]:
+    ) -> (
+        VersionedResult[StreamResponse] | VersionedResult[StreamResponse[FormattableT]]
+    ):
         """Stream the versioned function and return a wrapped StreamResponse."""
-        return self.stream.wrapped(*args, **kwargs)
+        return cast(
+            "VersionedResult[StreamResponse] | VersionedResult[StreamResponse[FormattableT]]",
+            self.stream.wrapped(*args, **kwargs),
+        )
 
 
 @dataclass(kw_only=True)
@@ -308,23 +346,65 @@ class VersionedAsyncCall(_BaseVersionedCall, Generic[P, FormattableT]):
         )
         self._compute_closure(self._call, self.call, self.stream)
 
+    @overload
+    async def __call__(
+        self: VersionedAsyncCall[P, None], *args: P.args, **kwargs: P.kwargs
+    ) -> AsyncResponse: ...
+
+    @overload
+    async def __call__(
+        self: VersionedAsyncCall[P, FormattableT], *args: P.args, **kwargs: P.kwargs
+    ) -> AsyncResponse[FormattableT]: ...
+
     async def __call__(
         self, *args: P.args, **kwargs: P.kwargs
     ) -> AsyncResponse | AsyncResponse[FormattableT]:
         """Call the versioned function and return AsyncResponse directly."""
         return await self.call(*args, **kwargs)
 
+    @overload
+    async def wrapped(
+        self: VersionedAsyncCall[P, None], *args: P.args, **kwargs: P.kwargs
+    ) -> AsyncVersionedResult[AsyncResponse]: ...
+
+    @overload
+    async def wrapped(
+        self: VersionedAsyncCall[P, FormattableT], *args: P.args, **kwargs: P.kwargs
+    ) -> AsyncVersionedResult[AsyncResponse[FormattableT]]: ...
+
     async def wrapped(
         self, *args: P.args, **kwargs: P.kwargs
-    ) -> AsyncVersionedResult[AsyncResponse | AsyncResponse[FormattableT]]:
+    ) -> (
+        AsyncVersionedResult[AsyncResponse]
+        | AsyncVersionedResult[AsyncResponse[FormattableT]]
+    ):
         """Call the versioned function and return a wrapped Response."""
-        return await self.call.wrapped(*args, **kwargs)
+        return cast(
+            "AsyncVersionedResult[AsyncResponse] | AsyncVersionedResult[AsyncResponse[FormattableT]]",
+            await self.call.wrapped(*args, **kwargs),
+        )
+
+    @overload
+    async def wrapped_stream(
+        self: VersionedAsyncCall[P, None], *args: P.args, **kwargs: P.kwargs
+    ) -> AsyncVersionedResult[AsyncStreamResponse]: ...
+
+    @overload
+    async def wrapped_stream(
+        self: VersionedAsyncCall[P, FormattableT], *args: P.args, **kwargs: P.kwargs
+    ) -> AsyncVersionedResult[AsyncStreamResponse[FormattableT]]: ...
 
     async def wrapped_stream(
         self, *args: P.args, **kwargs: P.kwargs
-    ) -> AsyncVersionedResult[AsyncStreamResponse | AsyncStreamResponse[FormattableT]]:
+    ) -> (
+        AsyncVersionedResult[AsyncStreamResponse]
+        | AsyncVersionedResult[AsyncStreamResponse[FormattableT]]
+    ):
         """Stream the versioned function and return a wrapped StreamResponse."""
-        return await self.stream.wrapped(*args, **kwargs)
+        return cast(
+            "AsyncVersionedResult[AsyncStreamResponse] | AsyncVersionedResult[AsyncStreamResponse[FormattableT]]",
+            await self.stream.wrapped(*args, **kwargs),
+        )
 
 
 @dataclass(kw_only=True)
@@ -381,18 +461,66 @@ class VersionedContextCall(_BaseVersionedCall, Generic[P, DepsT, FormattableT]):
         )
         self._compute_closure(self._call, self._call_versioned, self._stream_versioned)
 
+    @overload
+    def call(
+        self: VersionedContextCall[P, DepsT, None],
+        ctx: Context[DepsT],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> ContextResponse[DepsT, None]: ...
+
+    @overload
+    def call(
+        self: VersionedContextCall[P, DepsT, FormattableT],
+        ctx: Context[DepsT],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> ContextResponse[DepsT, FormattableT]: ...
+
     def call(
         self, ctx: Context[DepsT], *args: P.args, **kwargs: P.kwargs
     ) -> ContextResponse[DepsT, None] | ContextResponse[DepsT, FormattableT]:
         """Call the versioned function and return ContextResponse directly."""
         return self._call_versioned(ctx, *args, **kwargs)
 
+    @overload
+    def __call__(
+        self: VersionedContextCall[P, DepsT, None],
+        ctx: Context[DepsT],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> ContextResponse[DepsT, None]: ...
+
+    @overload
+    def __call__(
+        self: VersionedContextCall[P, DepsT, FormattableT],
+        ctx: Context[DepsT],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> ContextResponse[DepsT, FormattableT]: ...
+
     def __call__(
         self, ctx: Context[DepsT], *args: P.args, **kwargs: P.kwargs
     ) -> ContextResponse[DepsT, None] | ContextResponse[DepsT, FormattableT]:
         """Call the versioned function and return ContextResponse directly."""
         return self.call(ctx, *args, **kwargs)
 
+    @overload
+    def stream(
+        self: VersionedContextCall[P, DepsT, None],
+        ctx: Context[DepsT],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> ContextStreamResponse[DepsT, None]: ...
+
+    @overload
+    def stream(
+        self: VersionedContextCall[P, DepsT, FormattableT],
+        ctx: Context[DepsT],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> ContextStreamResponse[DepsT, FormattableT]: ...
+
     def stream(
         self, ctx: Context[DepsT], *args: P.args, **kwargs: P.kwargs
     ) -> (
@@ -401,6 +529,22 @@ class VersionedContextCall(_BaseVersionedCall, Generic[P, DepsT, FormattableT]):
         """Stream the versioned function and return a ContextStreamResponse."""
         return self._stream_versioned(ctx, *args, **kwargs)
 
+    @overload
+    def wrapped(
+        self: VersionedContextCall[P, DepsT, None],
+        ctx: Context[DepsT],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> VersionedResult[ContextResponse[DepsT, None]]: ...
+
+    @overload
+    def wrapped(
+        self: VersionedContextCall[P, DepsT, FormattableT],
+        ctx: Context[DepsT],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> VersionedResult[ContextResponse[DepsT, FormattableT]]: ...
+
     def wrapped(
         self, ctx: Context[DepsT], *args: P.args, **kwargs: P.kwargs
     ) -> VersionedResult[
@@ -409,6 +553,22 @@ class VersionedContextCall(_BaseVersionedCall, Generic[P, DepsT, FormattableT]):
         """Call the versioned function and return a wrapped Response."""
         return self._call_versioned.wrapped(ctx, *args, **kwargs)
 
+    @overload
+    def wrapped_stream(
+        self: VersionedContextCall[P, DepsT, None],
+        ctx: Context[DepsT],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> VersionedResult[ContextStreamResponse[DepsT, None]]: ...
+
+    @overload
+    def wrapped_stream(
+        self: VersionedContextCall[P, DepsT, FormattableT],
+        ctx: Context[DepsT],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> VersionedResult[ContextStreamResponse[DepsT, FormattableT]]: ...
+
     def wrapped_stream(
         self, ctx: Context[DepsT], *args: P.args, **kwargs: P.kwargs
     ) -> VersionedResult[
@@ -473,18 +633,66 @@ class VersionedAsyncContextCall(_BaseVersionedCall, Generic[P, DepsT, Formattabl
         )
         self._compute_closure(self._call, self._call_versioned, self._stream_versioned)
 
+    @overload
+    async def call(
+        self: VersionedAsyncContextCall[P, DepsT, None],
+        ctx: Context[DepsT],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> AsyncContextResponse[DepsT, None]: ...
+
+    @overload
+    async def call(
+        self: VersionedAsyncContextCall[P, DepsT, FormattableT],
+        ctx: Context[DepsT],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> AsyncContextResponse[DepsT, FormattableT]: ...
+
     async def call(
         self, ctx: Context[DepsT], *args: P.args, **kwargs: P.kwargs
     ) -> AsyncContextResponse[DepsT, None] | AsyncContextResponse[DepsT, FormattableT]:
         """Call the versioned function and return AsyncContextResponse directly."""
         return await self._call_versioned(ctx, *args, **kwargs)
 
+    @overload
+    async def __call__(
+        self: VersionedAsyncContextCall[P, DepsT, None],
+        ctx: Context[DepsT],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> AsyncContextResponse[DepsT, None]: ...
+
+    @overload
+    async def __call__(
+        self: VersionedAsyncContextCall[P, DepsT, FormattableT],
+        ctx: Context[DepsT],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> AsyncContextResponse[DepsT, FormattableT]: ...
+
     async def __call__(
         self, ctx: Context[DepsT], *args: P.args, **kwargs: P.kwargs
     ) -> AsyncContextResponse[DepsT, None] | AsyncContextResponse[DepsT, FormattableT]:
         """Call the versioned function and return AsyncContextResponse directly."""
         return await self.call(ctx, *args, **kwargs)
 
+    @overload
+    async def stream(
+        self: VersionedAsyncContextCall[P, DepsT, None],
+        ctx: Context[DepsT],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> AsyncContextStreamResponse[DepsT, None]: ...
+
+    @overload
+    async def stream(
+        self: VersionedAsyncContextCall[P, DepsT, FormattableT],
+        ctx: Context[DepsT],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> AsyncContextStreamResponse[DepsT, FormattableT]: ...
+
     async def stream(
         self, ctx: Context[DepsT], *args: P.args, **kwargs: P.kwargs
     ) -> (
@@ -494,6 +702,22 @@ class VersionedAsyncContextCall(_BaseVersionedCall, Generic[P, DepsT, Formattabl
         """Stream the versioned function and return an AsyncContextStreamResponse."""
         return await self._stream_versioned(ctx, *args, **kwargs)
 
+    @overload
+    async def wrapped(
+        self: VersionedAsyncContextCall[P, DepsT, None],
+        ctx: Context[DepsT],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> AsyncVersionedResult[AsyncContextResponse[DepsT, None]]: ...
+
+    @overload
+    async def wrapped(
+        self: VersionedAsyncContextCall[P, DepsT, FormattableT],
+        ctx: Context[DepsT],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> AsyncVersionedResult[AsyncContextResponse[DepsT, FormattableT]]: ...
+
     async def wrapped(
         self, ctx: Context[DepsT], *args: P.args, **kwargs: P.kwargs
     ) -> AsyncVersionedResult[
@@ -502,6 +726,22 @@ class VersionedAsyncContextCall(_BaseVersionedCall, Generic[P, DepsT, Formattabl
         """Call the versioned function and return a wrapped Response."""
         return await self._call_versioned.wrapped(ctx, *args, **kwargs)
 
+    @overload
+    async def wrapped_stream(
+        self: VersionedAsyncContextCall[P, DepsT, None],
+        ctx: Context[DepsT],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> AsyncVersionedResult[AsyncContextStreamResponse[DepsT, None]]: ...
+
+    @overload
+    async def wrapped_stream(
+        self: VersionedAsyncContextCall[P, DepsT, FormattableT],
+        ctx: Context[DepsT],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> AsyncVersionedResult[AsyncContextStreamResponse[DepsT, FormattableT]]: ...
+
     async def wrapped_stream(
         self, ctx: Context[DepsT], *args: P.args, **kwargs: P.kwargs
     ) -> AsyncVersionedResult[

mirascope/ops/_internal/versioned_functions.py

@@ -270,6 +270,7 @@ class VersionedFunction(_BaseVersionedFunction[P, R], BaseSyncTracedFunction[P,
                 signature=self.closure.signature,
                 signature_hash=self.closure.signature_hash,
                 name=self.name or self.closure.name,
+                language="python",
                 description=self.closure.docstring,
                 tags=list(self.tags) if self.tags else None,
                 metadata=cast(dict[str, str | None], self.metadata)
@@ -344,6 +345,7 @@ class AsyncVersionedFunction(
                 signature=self.closure.signature,
                 signature_hash=self.closure.signature_hash,
                 name=self.name or self.closure.name,
+                language="python",
                 description=self.closure.docstring,
                 tags=list(self.tags) if self.tags else None,
                 metadata=cast(dict[str, str | None], self.metadata)

{mirascope-2.1.0.dist-info → mirascope-2.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mirascope
-Version: 2.1.0
+Version: 2.2.0
 Summary: Every frontier LLM. One unified interface.
 Project-URL: Homepage, https://mirascope.com
 Project-URL: Documentation, https://mirascope.com/docs/mirascope/v2