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

mirascope/llm/providers/mlx/mlx.py

@@ -0,0 +1,242 @@
+from __future__ import annotations
+
+import asyncio
+import threading
+from collections.abc import Iterable, Sequence
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+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, FormatSpec, FormattableT
+from ...messages import AssistantMessage, Message, assistant
+from ...responses import AsyncChunkIterator, ChunkIterator, StreamResponseChunk
+from ...tools import AnyToolSchema, BaseToolkit
+from . import _utils
+from .encoding import BaseEncoder, TokenIds
+from .model_id import MLXModelId
+
+if TYPE_CHECKING:
+    from ...models import Params
+
+
+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: BaseToolkit[AnyToolSchema],
+        format: FormatSpec[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: BaseToolkit[AnyToolSchema],
+        format: FormatSpec[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: BaseToolkit[AnyToolSchema],
+        format: FormatSpec[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: BaseToolkit[AnyToolSchema],
+        format: FormatSpec[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
+"""

mirascope/llm/providers/mlx/provider.py

@@ -0,0 +1,416 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+from functools import cache, lru_cache
+from typing import TYPE_CHECKING, cast
+from typing_extensions import Unpack
+
+import mlx.nn as nn
+from mlx_lm import load as mlx_load
+from transformers import PreTrainedTokenizer
+
+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
+from . import _utils
+from .encoding import TransformersEncoder
+from .mlx import MLX
+from .model_id import MLXModelId
+
+if TYPE_CHECKING:
+    from ...models import Params
+
+
+@cache
+def _mlx_client_singleton() -> MLXProvider:
+    """Get or create the singleton MLX client instance."""
+    return MLXProvider()
+
+
+def client() -> MLXProvider:
+    """Get the MLX client singleton instance."""
+    return _mlx_client_singleton()
+
+
+@lru_cache(maxsize=16)
+def _get_mlx(model_id: MLXModelId) -> MLX:
+    model, tokenizer = cast(tuple[nn.Module, PreTrainedTokenizer], mlx_load(model_id))
+    encoder = TransformersEncoder(tokenizer)
+    return MLX(
+        model_id,
+        model,
+        tokenizer,
+        encoder,
+    )
+
+
+class MLXProvider(BaseProvider[None]):
+    """Client for interacting with MLX language models.
+
+    This client provides methods for generating responses from MLX models,
+    supporting both synchronous and asynchronous operations, as well as
+    streaming responses.
+    """
+
+    id = "mlx"
+    default_scope = "mlx-community/"
+    error_map = _utils.MLX_ERROR_MAP
+
+    def get_error_status(self, e: Exception) -> int | None:
+        """Extract HTTP status code from MLX exception.
+
+        MLX/HuggingFace Hub exceptions don't have status codes.
+        """
+        return None
+
+    def _call(
+        self,
+        *,
+        model_id: MLXModelId,
+        messages: Sequence[Message],
+        toolkit: Toolkit,
+        format: FormatSpec[FormattableT] | None = None,
+        **params: Unpack[Params],
+    ) -> Response | Response[FormattableT]:
+        """Generate an `llm.Response` using MLX model.
+
+        Args:
+            model_id: Model identifier to use.
+            messages: Messages to send to the LLM.
+            tools: Optional tools that the model may invoke.
+            format: Optional response format specifier.
+            **params: Additional parameters to configure output (e.g. temperature). See `llm.Params`.
+
+        Returns:
+            An `llm.Response` object containing the LLM-generated content.
+        """
+        mlx = _get_mlx(model_id)
+
+        input_messages, format, assistant_message, response = mlx.generate(
+            messages, toolkit, format, params
+        )
+
+        return Response(
+            raw=response,
+            provider_id="mlx",
+            model_id=model_id,
+            provider_model_name=model_id,
+            params=params,
+            tools=toolkit,
+            input_messages=input_messages,
+            assistant_message=assistant_message,
+            finish_reason=_utils.extract_finish_reason(response),
+            usage=_utils.extract_usage(response),
+            format=format,
+        )
+
+    def _context_call(
+        self,
+        *,
+        ctx: Context[DepsT],
+        model_id: MLXModelId,
+        messages: Sequence[Message],
+        toolkit: ContextToolkit[DepsT],
+        format: FormatSpec[FormattableT] | None = None,
+        **params: Unpack[Params],
+    ) -> ContextResponse[DepsT, None] | ContextResponse[DepsT, FormattableT]:
+        """Generate an `llm.ContextResponse` using MLX model.
+
+        Args:
+            ctx: Context object with dependencies for tools.
+            model_id: Model identifier to use.
+            messages: Messages to send to the LLM.
+            tools: Optional tools that the model may invoke.
+            format: Optional response format specifier.
+            **params: Additional parameters to configure output (e.g. temperature). See `llm.Params`.
+
+        Returns:
+            An `llm.ContextResponse` object containing the LLM-generated content.
+        """
+        mlx = _get_mlx(model_id)
+
+        input_messages, format, assistant_message, response = mlx.generate(
+            messages, toolkit, format, params
+        )
+
+        return ContextResponse(
+            raw=response,
+            provider_id="mlx",
+            model_id=model_id,
+            provider_model_name=model_id,
+            params=params,
+            tools=toolkit,
+            input_messages=input_messages,
+            assistant_message=assistant_message,
+            finish_reason=_utils.extract_finish_reason(response),
+            usage=_utils.extract_usage(response),
+            format=format,
+        )
+
+    async def _call_async(
+        self,
+        *,
+        model_id: MLXModelId,
+        messages: Sequence[Message],
+        toolkit: AsyncToolkit,
+        format: FormatSpec[FormattableT] | None = None,
+        **params: Unpack[Params],
+    ) -> AsyncResponse | AsyncResponse[FormattableT]:
+        """Generate an `llm.AsyncResponse` using MLX model by asynchronously calloing
+        `asycio.to_thread`.
+
+        Args:
+            model_id: Model identifier to use.
+            messages: Messages to send to the LLM.
+            tools: Optional tools that the model may invoke.
+            format: Optional response format specifier.
+            **params: Additional parameters to configure output (e.g. temperature). See `llm.Params`.
+
+        Returns:
+            An `llm.AsyncResponse` object containing the LLM-generated content.
+        """
+        mlx = _get_mlx(model_id)
+
+        (
+            input_messages,
+            format,
+            assistant_message,
+            response,
+        ) = await mlx.generate_async(messages, toolkit, format, params)
+
+        return AsyncResponse(
+            raw=response,
+            provider_id="mlx",
+            model_id=model_id,
+            provider_model_name=model_id,
+            params=params,
+            tools=toolkit,
+            input_messages=input_messages,
+            assistant_message=assistant_message,
+            finish_reason=_utils.extract_finish_reason(response),
+            usage=_utils.extract_usage(response),
+            format=format,
+        )
+
+    async def _context_call_async(
+        self,
+        *,
+        ctx: Context[DepsT],
+        model_id: MLXModelId,
+        messages: Sequence[Message],
+        toolkit: AsyncContextToolkit[DepsT],
+        format: FormatSpec[FormattableT] | None = None,
+        **params: Unpack[Params],
+    ) -> AsyncContextResponse[DepsT, None] | AsyncContextResponse[DepsT, FormattableT]:
+        """Generate an `llm.AsyncResponse` using MLX model by asynchronously calloing
+        `asycio.to_thread`.
+
+        Args:
+            ctx: Context object with dependencies for tools.
+            model_id: Model identifier to use.
+            messages: Messages to send to the LLM.
+            tools: Optional tools that the model may invoke.
+            format: Optional response format specifier.
+            **params: Additional parameters to configure output (e.g. temperature). See `llm.Params`.
+
+        Returns:
+            An `llm.AsyncContextResponse` object containing the LLM-generated content.
+        """
+        mlx = _get_mlx(model_id)
+
+        (
+            input_messages,
+            format,
+            assistant_message,
+            response,
+        ) = await mlx.generate_async(messages, toolkit, format, params)
+
+        return AsyncContextResponse(
+            raw=response,
+            provider_id="mlx",
+            model_id=model_id,
+            provider_model_name=model_id,
+            params=params,
+            tools=toolkit,
+            input_messages=input_messages,
+            assistant_message=assistant_message,
+            finish_reason=_utils.extract_finish_reason(response),
+            usage=_utils.extract_usage(response),
+            format=format,
+        )
+
+    def _stream(
+        self,
+        *,
+        model_id: MLXModelId,
+        messages: Sequence[Message],
+        toolkit: Toolkit,
+        format: FormatSpec[FormattableT] | None = None,
+        **params: Unpack[Params],
+    ) -> StreamResponse | StreamResponse[FormattableT]:
+        """Generate an `llm.StreamResponse` by synchronously streaming from MLX model output.
+
+        Args:
+            model_id: Model identifier to use.
+            messages: Messages to send to the LLM.
+            tools: Optional tools that the model may invoke.
+            format: Optional response format specifier.
+            **params: Additional parameters to configure output (e.g. temperature). See `llm.Params`.
+
+        Returns:
+            An `llm.StreamResponse` object for iterating over the LLM-generated content.
+        """
+        mlx = _get_mlx(model_id)
+
+        input_messages, format, chunk_iterator = mlx.stream(
+            messages, toolkit, format, params
+        )
+
+        return StreamResponse(
+            provider_id="mlx",
+            model_id=model_id,
+            provider_model_name=model_id,
+            params=params,
+            tools=toolkit,
+            input_messages=input_messages,
+            chunk_iterator=chunk_iterator,
+            format=format,
+        )
+
+    def _context_stream(
+        self,
+        *,
+        ctx: Context[DepsT],
+        model_id: MLXModelId,
+        messages: Sequence[Message],
+        toolkit: ContextToolkit[DepsT],
+        format: FormatSpec[FormattableT] | None = None,
+        **params: Unpack[Params],
+    ) -> ContextStreamResponse[DepsT] | ContextStreamResponse[DepsT, FormattableT]:
+        """Generate an `llm.ContextStreamResponse` by synchronously streaming from MLX model output.
+
+        Args:
+            ctx: Context object with dependencies for tools.
+            model_id: Model identifier to use.
+            messages: Messages to send to the LLM.
+            tools: Optional tools that the model may invoke.
+            format: Optional response format specifier.
+            **params: Additional parameters to configure output (e.g. temperature). See `llm.Params`.
+
+        Returns:
+            An `llm.ContextStreamResponse` object for iterating over the LLM-generated content.
+        """
+        mlx = _get_mlx(model_id)
+
+        input_messages, format, chunk_iterator = mlx.stream(
+            messages, toolkit, format, params
+        )
+
+        return ContextStreamResponse(
+            provider_id="mlx",
+            model_id=model_id,
+            provider_model_name=model_id,
+            params=params,
+            tools=toolkit,
+            input_messages=input_messages,
+            chunk_iterator=chunk_iterator,
+            format=format,
+        )
+
+    async def _stream_async(
+        self,
+        *,
+        model_id: MLXModelId,
+        messages: Sequence[Message],
+        toolkit: AsyncToolkit,
+        format: FormatSpec[FormattableT] | None = None,
+        **params: Unpack[Params],
+    ) -> AsyncStreamResponse | AsyncStreamResponse[FormattableT]:
+        """Generate an `llm.AsyncStreamResponse` by asynchronously streaming from MLX model output.
+
+        Args:
+            model_id: Model identifier to use.
+            messages: Messages to send to the LLM.
+            tools: Optional tools that the model may invoke.
+            format: Optional response format specifier.
+            **params: Additional parameters to configure output (e.g. temperature). See `llm.Params`.
+
+        Returns:
+            An `llm.AsyncStreamResponse` object for asynchronously iterating over the LLM-generated content.
+        """
+        mlx = _get_mlx(model_id)
+
+        input_messages, format, chunk_iterator = await mlx.stream_async(
+            messages, toolkit, format, params
+        )
+
+        return AsyncStreamResponse(
+            provider_id="mlx",
+            model_id=model_id,
+            provider_model_name=model_id,
+            params=params,
+            tools=toolkit,
+            input_messages=input_messages,
+            chunk_iterator=chunk_iterator,
+            format=format,
+        )
+
+    async def _context_stream_async(
+        self,
+        *,
+        ctx: Context[DepsT],
+        model_id: MLXModelId,
+        messages: Sequence[Message],
+        toolkit: AsyncContextToolkit[DepsT],
+        format: FormatSpec[FormattableT] | None = None,
+        **params: Unpack[Params],
+    ) -> (
+        AsyncContextStreamResponse[DepsT]
+        | AsyncContextStreamResponse[DepsT, FormattableT]
+    ):
+        """Generate an `llm.AsyncContextStreamResponse` by asynchronously streaming from MLX model output.
+
+        Args:
+            ctx: Context object with dependencies for tools.
+            model_id: Model identifier to use.
+            messages: Messages to send to the LLM.
+            tools: Optional tools that the model may invoke.
+            format: Optional response format specifier.
+            **params: Additional parameters to configure output (e.g. temperature). See `llm.Params`.
+
+        Returns:
+            An `llm.AsyncContextStreamResponse` object for asynchronously iterating over the LLM-generated content.
+        """
+        mlx = _get_mlx(model_id)
+
+        input_messages, format, chunk_iterator = await mlx.stream_async(
+            messages, toolkit, format, params
+        )
+
+        return AsyncContextStreamResponse(
+            provider_id="mlx",
+            model_id=model_id,
+            provider_model_name=model_id,
+            params=params,
+            tools=toolkit,
+            input_messages=input_messages,
+            chunk_iterator=chunk_iterator,
+            format=format,
+        )

mirascope/llm/providers/model_id.py

@@ -0,0 +1,16 @@
+from typing import TypeAlias
+
+from .anthropic import (
+    AnthropicModelId,
+)
+from .google import (
+    GoogleModelId,
+)
+from .mlx import (
+    MLXModelId,
+)
+from .openai import (
+    OpenAIModelId,
+)
+
+ModelId: TypeAlias = AnthropicModelId | GoogleModelId | OpenAIModelId | MLXModelId | str

mirascope/llm/providers/ollama/__init__.py

@@ -0,0 +1,7 @@
+"""Ollama provider implementation."""
+
+from .provider import OllamaProvider
+
+__all__ = [
+    "OllamaProvider",
+]