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

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

@@ -1,150 +0,0 @@
-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, OutputParser
-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]
-        | OutputParser[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

@@ -1,254 +0,0 @@
-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, FormattableT, OutputParser
-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: Sequence[AnyToolSchema] | BaseToolkit[AnyToolSchema] | None,
-        format: type[FormattableT]
-        | Format[FormattableT]
-        | OutputParser[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]
-        | OutputParser[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]
-        | OutputParser[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]
-        | OutputParser[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

@@ -1,17 +0,0 @@
-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
-"""