mirascope 1.25.6__py3-none-any.whl → 2.0.0__py3-none-any.whl

mirascope/llm/providers/mirascope/provider.py

@@ -0,0 +1,349 @@
+"""Mirascope Router provider that routes requests through the Mirascope Router API."""
+
+from __future__ import annotations
+
+import os
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
+from typing_extensions import Unpack
+
+from ...context import Context, DepsT
+from ...formatting import Format, FormattableT, OutputParser
+from ...messages import Message
+from ...responses import (
+    AsyncContextResponse,
+    AsyncContextStreamResponse,
+    AsyncResponse,
+    AsyncStreamResponse,
+    ContextResponse,
+    ContextStreamResponse,
+    Response,
+    StreamResponse,
+)
+from ...tools import (
+    AsyncContextTool,
+    AsyncContextToolkit,
+    AsyncTool,
+    AsyncToolkit,
+    ContextTool,
+    ContextToolkit,
+    Tool,
+    Toolkit,
+)
+from ..base import BaseProvider, Provider
+from . import _utils
+
+if TYPE_CHECKING:
+    from ...models import Params
+
+
+class MirascopeProvider(BaseProvider[None]):
+    """Provider that routes LLM requests through the Mirascope Router API.
+
+    The Mirascope Router provides a unified API for multiple LLM providers
+    (Anthropic, Google, OpenAI) with usage tracking and cost calculation.
+
+    This provider:
+    - Takes model IDs in the format "scope/model-name" (e.g., "openai/gpt-4")
+    - Routes requests to the Mirascope Router endpoint
+    - Delegates to the appropriate underlying provider (Anthropic, Google, or OpenAI)
+    - Uses MIRASCOPE_API_KEY for authentication
+
+    Environment Variables:
+        MIRASCOPE_API_KEY: Required API key for Mirascope Router authentication
+        MIRASCOPE_ROUTER_BASE_URL: Optional base URL override
+            (default: https://mirascope.com/router/v2)
+
+    Example:
+        ```python
+        import os
+        from mirascope import llm
+
+        os.environ["MIRASCOPE_API_KEY"] = "mk..."
+
+        # Register the Mirascope provider
+        llm.register_provider(
+            "mirascope",
+            scope=["anthropic/", "google/", "openai/"],
+        )
+
+        # Use with llm.call decorator
+        @llm.call("openai/gpt-4")
+        def recommend_book(genre: str):
+            return f"Recommend a {genre} book"
+
+        response = recommend_book("fantasy")
+        print(response.content)
+        ```
+    """
+
+    id = "mirascope"
+    default_scope = ["anthropic/", "google/", "openai/"]
+    error_map = {}
+    """Empty error map since MirascopeProvider delegates to underlying providers.
+
+    Error handling is performed by the underlying provider instances (Anthropic,
+    Google, OpenAI), which have their own error maps. Any exceptions that bubble
+    up from underlying providers are already converted to Mirascope exceptions.
+    """
+
+    def __init__(
+        self, *, api_key: str | None = None, base_url: str | None = None
+    ) -> None:
+        """Initialize the Mirascope provider.
+
+        Args:
+            api_key: Mirascope API key. If not provided, reads from MIRASCOPE_API_KEY
+                environment variable.
+            base_url: Optional base URL override for the Mirascope Router. If not
+                provided, reads from MIRASCOPE_ROUTER_BASE_URL environment variable
+                or defaults to https://mirascope.com/router/v2
+        """
+        api_key = api_key or os.environ.get("MIRASCOPE_API_KEY")
+        if not api_key:
+            raise ValueError(
+                "Mirascope API key not found. "
+                "Set MIRASCOPE_API_KEY environment variable or pass api_key parameter."
+            )
+
+        self.api_key = api_key
+        self.router_base_url = base_url or _utils.get_default_router_base_url()
+        self.client = None  # No single client; we create per-provider clients
+
+    def get_error_status(self, e: Exception) -> int | None:
+        """Extract HTTP status code from exception.
+
+        Since MirascopeProvider delegates to underlying providers, this method
+        is not used for direct error extraction. Underlying providers handle
+        their own status code extraction.
+
+        Args:
+            e: The exception to extract status code from.
+
+        Returns:
+            None, as status extraction is handled by underlying providers.
+        """
+        return None
+
+    def _get_underlying_provider(self, model_id: str) -> Provider:
+        """Get the underlying provider for a model ID.
+
+        Args:
+            model_id: Model identifier in format "scope/model-name"
+
+        Returns:
+            The appropriate cached provider instance (Anthropic, Google, or OpenAI)
+
+        Raises:
+            ValueError: If the model ID format is invalid or provider is unsupported
+        """
+        model_scope = _utils.extract_model_scope(model_id)
+        if not model_scope:
+            raise ValueError(
+                f"Invalid model ID format: {model_id}. "
+                f"Expected format 'scope/model-name' (e.g., 'openai/gpt-4')"
+            )
+
+        # Use the cached function to get/create the provider
+        return _utils.create_underlying_provider(
+            model_scope=model_scope,
+            api_key=self.api_key,
+            router_base_url=self.router_base_url,
+        )
+
+    def _call(
+        self,
+        *,
+        model_id: str,
+        messages: Sequence[Message],
+        tools: Sequence[Tool] | Toolkit | None = None,
+        format: type[FormattableT]
+        | Format[FormattableT]
+        | OutputParser[FormattableT]
+        | None = None,
+        **params: Unpack[Params],
+    ) -> Response | Response[FormattableT]:
+        """Generate an `llm.Response` by calling through the Mirascope Router."""
+        provider = self._get_underlying_provider(model_id)
+        return provider.call(
+            model_id=model_id,
+            messages=messages,
+            tools=tools,
+            format=format,
+            **params,
+        )
+
+    def _context_call(
+        self,
+        *,
+        ctx: Context[DepsT],
+        model_id: str,
+        messages: Sequence[Message],
+        tools: Sequence[Tool | ContextTool[DepsT]]
+        | ContextToolkit[DepsT]
+        | None = None,
+        format: type[FormattableT]
+        | Format[FormattableT]
+        | OutputParser[FormattableT]
+        | None = None,
+        **params: Unpack[Params],
+    ) -> ContextResponse[DepsT, None] | ContextResponse[DepsT, FormattableT]:
+        """Generate an `llm.ContextResponse` by calling through the Mirascope Router."""
+        provider = self._get_underlying_provider(model_id)
+        return provider.context_call(
+            ctx=ctx,
+            model_id=model_id,
+            messages=messages,
+            tools=tools,
+            format=format,
+            **params,
+        )
+
+    async def _call_async(
+        self,
+        *,
+        model_id: str,
+        messages: Sequence[Message],
+        tools: Sequence[AsyncTool] | AsyncToolkit | None = None,
+        format: type[FormattableT]
+        | Format[FormattableT]
+        | OutputParser[FormattableT]
+        | None = None,
+        **params: Unpack[Params],
+    ) -> AsyncResponse | AsyncResponse[FormattableT]:
+        """Generate an `llm.AsyncResponse` by calling through the Mirascope Router."""
+        provider = self._get_underlying_provider(model_id)
+        return await provider.call_async(
+            model_id=model_id,
+            messages=messages,
+            tools=tools,
+            format=format,
+            **params,
+        )
+
+    async def _context_call_async(
+        self,
+        *,
+        ctx: Context[DepsT],
+        model_id: str,
+        messages: Sequence[Message],
+        tools: Sequence[AsyncTool | AsyncContextTool[DepsT]]
+        | AsyncContextToolkit[DepsT]
+        | None = None,
+        format: type[FormattableT]
+        | Format[FormattableT]
+        | OutputParser[FormattableT]
+        | None = None,
+        **params: Unpack[Params],
+    ) -> AsyncContextResponse[DepsT, None] | AsyncContextResponse[DepsT, FormattableT]:
+        """Generate an `llm.AsyncContextResponse` by calling through the Mirascope Router."""
+        provider = self._get_underlying_provider(model_id)
+        return await provider.context_call_async(
+            ctx=ctx,
+            model_id=model_id,
+            messages=messages,
+            tools=tools,
+            format=format,
+            **params,
+        )
+
+    def _stream(
+        self,
+        *,
+        model_id: str,
+        messages: Sequence[Message],
+        tools: Sequence[Tool] | Toolkit | None = None,
+        format: type[FormattableT]
+        | Format[FormattableT]
+        | OutputParser[FormattableT]
+        | None = None,
+        **params: Unpack[Params],
+    ) -> StreamResponse | StreamResponse[FormattableT]:
+        """Stream an `llm.StreamResponse` by calling through the Mirascope Router."""
+        provider = self._get_underlying_provider(model_id)
+        return provider.stream(
+            model_id=model_id,
+            messages=messages,
+            tools=tools,
+            format=format,
+            **params,
+        )
+
+    def _context_stream(
+        self,
+        *,
+        ctx: Context[DepsT],
+        model_id: str,
+        messages: Sequence[Message],
+        tools: Sequence[Tool | ContextTool[DepsT]]
+        | ContextToolkit[DepsT]
+        | None = None,
+        format: type[FormattableT]
+        | Format[FormattableT]
+        | OutputParser[FormattableT]
+        | None = None,
+        **params: Unpack[Params],
+    ) -> (
+        ContextStreamResponse[DepsT, None] | ContextStreamResponse[DepsT, FormattableT]
+    ):
+        """Stream an `llm.ContextStreamResponse` by calling through the Mirascope Router."""
+        provider = self._get_underlying_provider(model_id)
+        return provider.context_stream(
+            ctx=ctx,
+            model_id=model_id,
+            messages=messages,
+            tools=tools,
+            format=format,
+            **params,
+        )
+
+    async def _stream_async(
+        self,
+        *,
+        model_id: str,
+        messages: Sequence[Message],
+        tools: Sequence[AsyncTool] | AsyncToolkit | None = None,
+        format: type[FormattableT]
+        | Format[FormattableT]
+        | OutputParser[FormattableT]
+        | None = None,
+        **params: Unpack[Params],
+    ) -> AsyncStreamResponse | AsyncStreamResponse[FormattableT]:
+        """Stream an `llm.AsyncStreamResponse` by calling through the Mirascope Router."""
+        provider = self._get_underlying_provider(model_id)
+        return await provider.stream_async(
+            model_id=model_id,
+            messages=messages,
+            tools=tools,
+            format=format,
+            **params,
+        )
+
+    async def _context_stream_async(
+        self,
+        *,
+        ctx: Context[DepsT],
+        model_id: str,
+        messages: Sequence[Message],
+        tools: Sequence[AsyncTool | AsyncContextTool[DepsT]]
+        | AsyncContextToolkit[DepsT]
+        | None = None,
+        format: type[FormattableT]
+        | Format[FormattableT]
+        | OutputParser[FormattableT]
+        | None = None,
+        **params: Unpack[Params],
+    ) -> (
+        AsyncContextStreamResponse[DepsT, None]
+        | AsyncContextStreamResponse[DepsT, FormattableT]
+    ):
+        """Stream an `llm.AsyncContextStreamResponse` by calling through the Mirascope Router."""
+        provider = self._get_underlying_provider(model_id)
+        return await provider.context_stream_async(
+            ctx=ctx,
+            model_id=model_id,
+            messages=messages,
+            tools=tools,
+            format=format,
+            **params,
+        )

mirascope/llm/providers/mlx/__init__.py

@@ -0,0 +1,9 @@
+"""MLX client implementation."""
+
+from .model_id import MLXModelId
+from .provider import MLXProvider
+
+__all__ = [
+    "MLXModelId",
+    "MLXProvider",
+]

mirascope/llm/providers/mlx/_utils.py

@@ -0,0 +1,141 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import TYPE_CHECKING, TypeAlias, TypedDict
+
+import mlx.core as mx
+from huggingface_hub.errors import LocalEntryNotFoundError
+from mlx_lm.generate import GenerationResponse
+from mlx_lm.sample_utils import make_sampler
+
+from ...exceptions import NotFoundError
+from ...responses import FinishReason, Usage
+from ..base import ProviderErrorMap, _utils as _base_utils
+
+if TYPE_CHECKING:
+    from ...models import Params
+
+Sampler: TypeAlias = Callable[[mx.array], mx.array]
+
+# Error mapping for MLX provider
+MLX_ERROR_MAP: ProviderErrorMap = {
+    LocalEntryNotFoundError: NotFoundError,
+}
+
+
+class MakeSamplerKwargs(TypedDict, total=False):
+    """Keyword arguments to be used for `mlx_lm`-s `make_sampler` function.
+
+    Some of these settings are directly match the generic client parameters
+    as defined in the `Params` class. See mirascope.llm.providers.Params for
+    more details.
+    """
+
+    temp: float
+    "The temperature for sampling, if 0 the argmax is used."
+
+    top_p: float
+    "Nulceus sampling, higher means model considers more less likely words."
+
+    min_p: float
+    """The minimum value (scaled by the top token's probability) that a token
+    probability must have to be considered."""
+
+    min_tokens_to_keep: int
+    "Minimum number of tokens that cannot be filtered by min_p sampling."
+
+    top_k: int
+    "The top k tokens ranked by probability to constrain the sampling to."
+
+    xtc_probability: float
+    "The probability of applying XTC sampling."
+
+    xtc_threshold: float
+    "The threshold the probs need to reach for being sampled."
+
+    xtc_special_tokens: list[int]
+    "List of special tokens IDs to be excluded from XTC sampling."
+
+
+class StreamGenerateKwargs(TypedDict, total=False):
+    """Keyword arguments for the `mlx-lm.stream_generate` function."""
+
+    max_tokens: int
+    "The maximum number of tokens to generate."
+
+    sampler: Sampler
+    "A sampler for sampling token from a vector of logits."
+
+
+def encode_params(params: Params) -> tuple[int | None, StreamGenerateKwargs]:
+    """Convert generic params to mlx-lm stream_generate kwargs.
+
+    Args:
+        params: The generic parameters.
+
+    Returns:
+        The mlx-lm specific stream_generate keyword arguments.
+    """
+    kwargs: StreamGenerateKwargs = {}
+
+    with _base_utils.ensure_all_params_accessed(
+        params=params,
+        provider_id="mlx",
+        unsupported_params=["stop_sequences", "thinking"],
+    ) as param_accessor:
+        if param_accessor.max_tokens is not None:
+            kwargs["max_tokens"] = param_accessor.max_tokens
+        else:
+            kwargs["max_tokens"] = -1
+
+        sampler_kwargs = MakeSamplerKwargs({})
+        if param_accessor.temperature is not None:
+            sampler_kwargs["temp"] = param_accessor.temperature
+        if param_accessor.top_k is not None:
+            sampler_kwargs["top_k"] = param_accessor.top_k
+        if param_accessor.top_p is not None:
+            sampler_kwargs["top_p"] = param_accessor.top_p
+
+        kwargs["sampler"] = make_sampler(**sampler_kwargs)
+
+        return param_accessor.seed, kwargs
+
+
+def extract_finish_reason(response: GenerationResponse | None) -> FinishReason | None:
+    """Extract the finish reason from an MLX generation response.
+
+    Args:
+        response: The MLX generation response to extract from.
+
+    Returns:
+        The normalized finish reason, or None if not applicable.
+    """
+    if response is None:
+        return None
+
+    if response.finish_reason == "length":
+        return FinishReason.MAX_TOKENS
+
+    return None
+
+
+def extract_usage(response: GenerationResponse | None) -> Usage | None:
+    """Extract usage information from an MLX generation response.
+
+    Args:
+        response: The MLX generation response to extract from.
+
+    Returns:
+        The Usage object with token counts, or None if not applicable.
+    """
+    if response is None:
+        return None
+
+    return Usage(
+        input_tokens=response.prompt_tokens,
+        output_tokens=response.generation_tokens,
+        cache_read_tokens=0,
+        cache_write_tokens=0,
+        reasoning_tokens=0,
+        raw=response,
+    )

mirascope/llm/providers/mlx/encoding/__init__.py

@@ -0,0 +1,8 @@
+from .base import BaseEncoder, TokenIds
+from .transformers import TransformersEncoder
+
+__all__ = [
+    "BaseEncoder",
+    "TokenIds",
+    "TransformersEncoder",
+]

mirascope/llm/providers/mlx/encoding/base.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+import abc
+from collections.abc import Iterable, Sequence
+from typing import TypeAlias
+
+from mlx_lm.generate import GenerationResponse
+
+from ....formatting import Format, FormattableT, OutputParser
+from ....messages import AssistantContent, Message
+from ....responses import ChunkIterator
+from ....tools import AnyToolSchema, BaseToolkit
+
+TokenIds: TypeAlias = list[int]
+
+
+class BaseEncoder(abc.ABC):
+    """Abstract base class for Mirascope <> MLX encoding and decoding."""
+
+    @abc.abstractmethod
+    def encode_request(
+        self,
+        messages: Sequence[Message],
+        tools: Sequence[AnyToolSchema] | BaseToolkit[AnyToolSchema] | None,
+        format: type[FormattableT]
+        | Format[FormattableT]
+        | OutputParser[FormattableT]
+        | None,
+    ) -> tuple[Sequence[Message], Format[FormattableT] | None, TokenIds]:
+        """Encode the request messages into a format suitable for the model.
+
+        Args:
+            messages: The sequence of messages to encode.
+            tools: Optional sequence of tool schemas or toolkit for the model.
+            format: Optional format specification for structured outputs.
+
+        Returns:
+            A tuple containing:
+                - The processed messages
+                - The format specification (if provided)
+                - The encoded prompt as token IDs
+        """
+
+        ...
+
+    @abc.abstractmethod
+    def decode_response(
+        self, stream: Iterable[GenerationResponse]
+    ) -> tuple[AssistantContent, GenerationResponse | None]:
+        """Decode a stream of MLX generation responses into assistant content.
+
+        Args:
+            stream: An iterable of MLX generation responses.
+
+        Returns:
+            A tuple containing:
+                - The decoded assistant content
+                - The final generation response (if available)
+        """
+        ...
+
+    @abc.abstractmethod
+    def decode_stream(self, stream: Iterable[GenerationResponse]) -> ChunkIterator:
+        """Decode a stream of MLX generation responses into an iterable of chunks.
+
+        Args:
+            stream: An iterable of MLX generation responses.
+
+        Returns:
+            A ChunkIterator yielding content chunks for streaming responses.
+        """
+        ...