mirascope 2.0.0a2__py3-none-any.whl → 2.0.0a3__py3-none-any.whl

mirascope/llm/providers/mlx/__init__.py

@@ -0,0 +1,24 @@
+"""MLX client implementation."""
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .model_id import MLXModelId
+    from .provider import MLXProvider
+else:
+    try:
+        from .model_id import MLXModelId
+        from .provider import MLXProvider
+    except ImportError:  # pragma: no cover
+        from .._missing_import_stubs import (
+            create_import_error_stub,
+            create_provider_stub,
+        )
+
+        MLXProvider = create_provider_stub("mlx", "MLXProvider")
+        MLXModelId = str
+
+__all__ = [
+    "MLXModelId",
+    "MLXProvider",
+]

mirascope/llm/providers/mlx/_utils.py

@@ -0,0 +1,107 @@
+from collections.abc import Callable
+from typing import TypeAlias, TypedDict
+
+import mlx.core as mx
+from mlx_lm.generate import GenerationResponse
+from mlx_lm.sample_utils import make_sampler
+
+from ...responses import FinishReason
+from ..base import Params, _utils as _base_utils
+
+Sampler: TypeAlias = Callable[[mx.array], mx.array]
+
+
+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", "encode_thoughts_as_text"],
+    ) 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

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, 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: Sequence[AnyToolSchema] | BaseToolkit[AnyToolSchema] | None,
+        format: type[FormattableT] | Format[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,131 @@
+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, FormattableT
+from ....messages import AssistantContent, Message
+from ....responses import ChunkIterator, FinishReasonChunk, RawStreamEventChunk
+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: Sequence[AnyToolSchema] | BaseToolkit[AnyToolSchema] | None,
+        format: type[FormattableT] | Format[FormattableT] | None,
+    ) -> tuple[Sequence[Message], Format[FormattableT] | None, TokenIds]:
+        """Encode a request into a format suitable for the model."""
+        tool_schemas = tools.tools if isinstance(tools, BaseToolkit) else tools or []
+        if len(tool_schemas) > 0:
+            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()

mirascope/llm/providers/mlx/mlx.py

@@ -0,0 +1,237 @@
+import asyncio
+import threading
+from collections.abc import Iterable, Sequence
+from dataclasses import dataclass, field
+from typing_extensions import Unpack
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx_lm import stream_generate  # type: ignore[reportPrivateImportUsage]
+from mlx_lm.generate import GenerationResponse
+from transformers import PreTrainedTokenizer
+
+from ...formatting import Format, FormattableT
+from ...messages import AssistantMessage, Message, assistant
+from ...responses import AsyncChunkIterator, ChunkIterator, StreamResponseChunk
+from ...tools import AnyToolSchema, BaseToolkit
+from ..base import Params
+from . import _utils
+from .encoding import BaseEncoder, TokenIds
+from .model_id import MLXModelId
+
+
+def _consume_sync_stream_into_queue(
+    generation_stream: ChunkIterator,
+    loop: asyncio.AbstractEventLoop,
+    queue: asyncio.Queue[StreamResponseChunk | Exception | None],
+) -> None:
+    """Consume a synchronous stream and put chunks into an async queue.
+
+    Args:
+        sync_stream: The synchronous chunk iterator to consume.
+        loop: The event loop for scheduling queue operations.
+        queue: The async queue to put chunks into.
+    """
+    try:
+        for response in generation_stream:
+            asyncio.run_coroutine_threadsafe(queue.put(response), loop)
+    except Exception as e:
+        asyncio.run_coroutine_threadsafe(queue.put(e), loop)
+
+    asyncio.run_coroutine_threadsafe(queue.put(None), loop)
+
+
+@dataclass(frozen=True)
+class MLX:
+    """MLX model wrapper for synchronous and asynchronous generation.
+
+    Args:
+        model_id: The MLX model identifier.
+        model: The underlying MLX model.
+        tokenizer: The tokenizer for the model.
+        encoder: The encoder for prompts and responses.
+    """
+
+    model_id: MLXModelId
+    """The MLX model identifier."""
+
+    model: nn.Module
+    """The underlying MLX model."""
+
+    tokenizer: PreTrainedTokenizer
+    """The tokenizer for the model."""
+
+    encoder: BaseEncoder
+    """The encoder for prompts and responses."""
+
+    _lock: threading.Lock = field(default_factory=threading.Lock)
+    """The lock for thread-safety."""
+
+    def _stream_generate(
+        self,
+        prompt: TokenIds,
+        seed: int | None,
+        **kwargs: Unpack[_utils.StreamGenerateKwargs],
+    ) -> Iterable[GenerationResponse]:
+        """Generator that streams generation responses.
+
+        Using this generator instead of calling stream_generate directly ensures
+        thread-safety when using the model in a multi-threaded context.
+        """
+        with self._lock:
+            if seed is not None:
+                mx.random.seed(seed)
+
+            return stream_generate(
+                self.model,
+                self.tokenizer,
+                prompt,
+                **kwargs,
+            )
+
+    async def _stream_generate_async(
+        self,
+        prompt: TokenIds,
+        seed: int | None,
+        **kwargs: Unpack[_utils.StreamGenerateKwargs],
+    ) -> AsyncChunkIterator:
+        """Async generator that streams generation responses.
+
+        Note that, while stream_generate returns an iterable of GenerationResponse,
+        here we return an `AsyncChunkIterator`, in order to avoid having to implement
+        both synchronous and asynchronous versions of BaseEncoder.decode_stream.
+        This makes sense as in this case, there is nothing to gain from consuming the
+        generation asyncnronously.
+        """
+        loop = asyncio.get_running_loop()
+        generation_queue: asyncio.Queue[StreamResponseChunk | Exception | None] = (
+            asyncio.Queue()
+        )
+
+        sync_stream = self.encoder.decode_stream(
+            self._stream_generate(
+                prompt,
+                seed,
+                **kwargs,
+            )
+        )
+
+        consume_task = asyncio.create_task(
+            asyncio.to_thread(
+                _consume_sync_stream_into_queue, sync_stream, loop, generation_queue
+            ),
+        )
+        while item := await generation_queue.get():
+            if isinstance(item, Exception):
+                raise item
+
+            yield item
+
+        await consume_task
+
+    def stream(
+        self,
+        messages: Sequence[Message],
+        tools: Sequence[AnyToolSchema] | BaseToolkit[AnyToolSchema] | None,
+        format: type[FormattableT] | Format[FormattableT] | None,
+        params: Params,
+    ) -> tuple[Sequence[Message], Format[FormattableT] | None, ChunkIterator]:
+        """Stream response chunks synchronously.
+
+        Args:
+            messages: The input messages.
+            tools: Optional tools for the model.
+            format: Optional response format.
+
+        Returns:
+            Tuple of messages, format, and chunk iterator.
+        """
+        messages, format, prompt = self.encoder.encode_request(messages, tools, format)
+        seed, kwargs = _utils.encode_params(params)
+
+        stream = self._stream_generate(prompt, seed, **kwargs)
+        return messages, format, self.encoder.decode_stream(stream)
+
+    async def stream_async(
+        self,
+        messages: Sequence[Message],
+        tools: Sequence[AnyToolSchema] | BaseToolkit[AnyToolSchema] | None,
+        format: type[FormattableT] | Format[FormattableT] | None,
+        params: Params,
+    ) -> tuple[Sequence[Message], Format[FormattableT] | None, AsyncChunkIterator]:
+        """Stream response chunks asynchronously.
+
+        Args:
+            messages: The input messages.
+            tools: Optional tools for the model.
+            format: Optional response format.
+        Returns:
+            Tuple of messages, format, and async chunk iterator.
+        """
+        messages, format, prompt = await asyncio.to_thread(
+            self.encoder.encode_request, messages, tools, format
+        )
+        seed, kwargs = _utils.encode_params(params)
+
+        chunk_iterator = self._stream_generate_async(prompt, seed, **kwargs)
+        return messages, format, chunk_iterator
+
+    def generate(
+        self,
+        messages: Sequence[Message],
+        tools: Sequence[AnyToolSchema] | BaseToolkit[AnyToolSchema] | None,
+        format: type[FormattableT] | Format[FormattableT] | None,
+        params: Params,
+    ) -> tuple[
+        Sequence[Message],
+        Format[FormattableT] | None,
+        AssistantMessage,
+        GenerationResponse | None,
+    ]:
+        """Generate a response synchronously.
+
+        Args:
+            messages: The input messages.
+            tools: Optional tools for the model.
+            format: Optional response format.
+            params: Generation parameters.
+        Returns:
+            Tuple of messages, format, assistant message, and last generation response.
+        """
+        messages, format, prompt = self.encoder.encode_request(messages, tools, format)
+        seed, kwargs = _utils.encode_params(params)
+
+        stream = self._stream_generate(prompt, seed, **kwargs)
+        assistant_content, last_response = self.encoder.decode_response(stream)
+        assistant_message = assistant(
+            content=assistant_content,
+            model_id=self.model_id,
+            provider_id="mlx",
+            raw_message=None,
+            name=None,
+        )
+        return messages, format, assistant_message, last_response
+
+    async def generate_async(
+        self,
+        messages: Sequence[Message],
+        tools: Sequence[AnyToolSchema] | BaseToolkit[AnyToolSchema] | None,
+        format: type[FormattableT] | Format[FormattableT] | None,
+        params: Params,
+    ) -> tuple[
+        Sequence[Message],
+        Format[FormattableT] | None,
+        AssistantMessage,
+        GenerationResponse | None,
+    ]:
+        """Generate a response asynchronously.
+
+        Args:
+            messages: The input messages.
+            tools: Optional tools for the model.
+            format: Optional response format.
+            params: Generation parameters.
+        Returns:
+            Tuple of messages, format, assistant message, and last generation response.
+        """
+        return await asyncio.to_thread(self.generate, messages, tools, format, params)

mirascope/llm/providers/mlx/model_id.py

@@ -0,0 +1,17 @@
+from typing import TypeAlias
+
+# TODO: Add more explicit literals
+# TODO: Ensure automatic model downloads are supported.
+# TODO: Ensure instructions are clear for examples that run as copied
+MLXModelId: TypeAlias = str
+"""The identifier of the MLX model to be loaded by the MLX client.
+
+An MLX model identifier might be a local path to a model's file, or a huggingface
+repository such as:
+ - "mlx-community/Qwen3-8B-4bit-DWQ-053125"
+ - "mlx-community/gpt-oss-20b-MXFP4-Q8"
+
+For more details, see:
+ - https://github.com/ml-explore/mlx-lm/?tab=readme-ov-file#supported-models
+ - https://huggingface.co/mlx-community
+"""