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

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,147 @@
+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,
+    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: 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()
+
+        # 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,
+            )

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
+"""