mirascope 1.0.5__py3-none-any.whl → 2.1.1__py3-none-any.whl

mirascope/llm/providers/mirascope/provider.py

@@ -0,0 +1,313 @@
+"""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 FormatSpec, FormattableT
+from ...messages import Message
+from ...responses import (
+    AsyncContextResponse,
+    AsyncContextStreamResponse,
+    AsyncResponse,
+    AsyncStreamResponse,
+    ContextResponse,
+    ContextStreamResponse,
+    Response,
+    StreamResponse,
+)
+from ...tools import (
+    AsyncContextToolkit,
+    AsyncToolkit,
+    ContextToolkit,
+    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],
+        toolkit: Toolkit,
+        format: FormatSpec[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,
+            toolkit=toolkit,
+            format=format,
+            **params,
+        )
+
+    def _context_call(
+        self,
+        *,
+        ctx: Context[DepsT],
+        model_id: str,
+        messages: Sequence[Message],
+        toolkit: ContextToolkit[DepsT],
+        format: FormatSpec[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,
+            toolkit=toolkit,
+            format=format,
+            **params,
+        )
+
+    async def _call_async(
+        self,
+        *,
+        model_id: str,
+        messages: Sequence[Message],
+        toolkit: AsyncToolkit,
+        format: FormatSpec[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,
+            toolkit=toolkit,
+            format=format,
+            **params,
+        )
+
+    async def _context_call_async(
+        self,
+        *,
+        ctx: Context[DepsT],
+        model_id: str,
+        messages: Sequence[Message],
+        toolkit: AsyncContextToolkit[DepsT],
+        format: FormatSpec[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,
+            toolkit=toolkit,
+            format=format,
+            **params,
+        )
+
+    def _stream(
+        self,
+        *,
+        model_id: str,
+        messages: Sequence[Message],
+        toolkit: Toolkit,
+        format: FormatSpec[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,
+            toolkit=toolkit,
+            format=format,
+            **params,
+        )
+
+    def _context_stream(
+        self,
+        *,
+        ctx: Context[DepsT],
+        model_id: str,
+        messages: Sequence[Message],
+        toolkit: ContextToolkit[DepsT],
+        format: FormatSpec[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,
+            toolkit=toolkit,
+            format=format,
+            **params,
+        )
+
+    async def _stream_async(
+        self,
+        *,
+        model_id: str,
+        messages: Sequence[Message],
+        toolkit: AsyncToolkit,
+        format: FormatSpec[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,
+            toolkit=toolkit,
+            format=format,
+            **params,
+        )
+
+    async def _context_stream_async(
+        self,
+        *,
+        ctx: Context[DepsT],
+        model_id: str,
+        messages: Sequence[Message],
+        toolkit: AsyncContextToolkit[DepsT],
+        format: FormatSpec[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,
+            toolkit=toolkit,
+            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,69 @@
+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, FormatSpec, FormattableT
+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: BaseToolkit[AnyToolSchema],
+        format: FormatSpec[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.
+        """
+        ...

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

@@ -0,0 +1,146 @@
+import io
+from collections.abc import Iterable, Sequence
+from dataclasses import dataclass
+from typing import Literal, cast
+from typing_extensions import TypedDict
+
+from mlx_lm.generate import GenerationResponse
+from transformers import PreTrainedTokenizer
+
+from ....content import ContentPart, TextChunk, TextEndChunk, TextStartChunk
+from ....formatting import Format, FormatSpec, FormattableT
+from ....messages import AssistantContent, Message
+from ....responses import (
+    ChunkIterator,
+    FinishReasonChunk,
+    RawStreamEventChunk,
+    UsageDeltaChunk,
+)
+from ....tools import AnyToolSchema, BaseToolkit
+from .. import _utils
+from .base import BaseEncoder, TokenIds
+
+HFRole = Literal["system", "user", "assistant"] | str
+
+
+class TransformersMessage(TypedDict):
+    """Message in Transformers format."""
+
+    role: HFRole
+    content: str
+
+
+def _encode_content(content: Sequence[ContentPart]) -> str:
+    """Encode content parts into a string.
+
+    Args:
+        content: The sequence of content parts to encode.
+
+    Returns:
+        The encoded content as a string.
+
+    Raises:
+        NotImplementedError: If content contains non-text parts.
+    """
+    if len(content) == 1 and content[0].type == "text":
+        return content[0].text
+
+    raise NotImplementedError("Only text content is supported in this example.")
+
+
+def _encode_message(message: Message) -> TransformersMessage:
+    """Encode a Mirascope message into Transformers format.
+
+    Args:
+        message: The message to encode.
+
+    Returns:
+        The encoded message in Transformers format.
+
+    Raises:
+        ValueError: If the message role is not supported.
+    """
+    if message.role == "system":
+        return TransformersMessage(role="system", content=message.content.text)
+    elif message.role == "assistant" or message.role == "user":
+        return TransformersMessage(
+            role=message.role, content=_encode_content(message.content)
+        )
+    else:
+        raise ValueError(f"Unsupported message type: {type(message)}")
+
+
+@dataclass(frozen=True)
+class TransformersEncoder(BaseEncoder):
+    """Encoder for Transformers models."""
+
+    tokenizer: PreTrainedTokenizer
+    """The tokenizer to use for encoding."""
+
+    def encode_request(
+        self,
+        messages: Sequence[Message],
+        tools: BaseToolkit[AnyToolSchema],
+        format: FormatSpec[FormattableT] | None,
+    ) -> tuple[Sequence[Message], Format[FormattableT] | None, TokenIds]:
+        """Encode a request into a format suitable for the model."""
+        if tools.tools:
+            raise NotImplementedError("Tool usage is not supported.")
+        if format is not None:
+            raise NotImplementedError("Formatting is not supported.")
+
+        hf_messages: list[TransformersMessage] = [
+            _encode_message(msg) for msg in messages
+        ]
+        prompt_text = cast(
+            str,
+            self.tokenizer.apply_chat_template(  # pyright: ignore[reportUnknownMemberType]
+                cast(list[dict[str, str]], hf_messages),
+                tokenize=False,
+                add_generation_prompt=True,
+            ),
+        )
+        return (
+            messages,
+            format,
+            self.tokenizer.encode(prompt_text, add_special_tokens=False),  # pyright: ignore[reportUnknownMemberType]
+        )
+
+    def decode_response(
+        self, stream: Iterable[GenerationResponse]
+    ) -> tuple[AssistantContent, GenerationResponse | None]:
+        """Decode a response into a format suitable for the model."""
+        with io.StringIO() as buffer:
+            last_response: GenerationResponse | None = None
+            for response in stream:
+                buffer.write(response.text)
+                last_response = response
+
+            return buffer.getvalue(), last_response
+
+    def decode_stream(self, stream: Iterable[GenerationResponse]) -> ChunkIterator:
+        """Decode a stream of responses into a format suitable for the model."""
+        yield TextStartChunk()
+
+        response: GenerationResponse | None = None
+        for response in stream:
+            yield RawStreamEventChunk(raw_stream_event=response)
+            yield TextChunk(delta=response.text)
+
+        assert response is not None
+        finish_reason = _utils.extract_finish_reason(response)
+        if finish_reason is not None:
+            yield FinishReasonChunk(finish_reason=finish_reason)
+        else:
+            yield TextEndChunk()
+
+        # Emit usage delta if available
+        usage = _utils.extract_usage(response)
+        if usage:
+            yield UsageDeltaChunk(
+                input_tokens=usage.input_tokens,
+                output_tokens=usage.output_tokens,
+                cache_read_tokens=usage.cache_read_tokens,
+                cache_write_tokens=usage.cache_write_tokens,
+                reasoning_tokens=usage.reasoning_tokens,
+            )