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

mirascope/llm/providers/mlx/provider.py

@@ -0,0 +1,415 @@
+from collections.abc import Sequence
+from functools import cache, lru_cache
+from typing import 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 Format, FormattableT
+from ...messages import Message
+from ...responses import (
+    AsyncContextResponse,
+    AsyncContextStreamResponse,
+    AsyncResponse,
+    AsyncStreamResponse,
+    ContextResponse,
+    ContextStreamResponse,
+    Response,
+    StreamResponse,
+)
+from ...tools import (
+    AsyncContextTool,
+    AsyncContextToolkit,
+    AsyncTool,
+    AsyncToolkit,
+    ContextTool,
+    ContextToolkit,
+    Tool,
+    Toolkit,
+)
+from ..base import BaseProvider, Params
+from . import _utils
+from .encoding import TransformersEncoder
+from .mlx import MLX
+from .model_id import MLXModelId
+
+
+@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/"
+
+    def _call(
+        self,
+        *,
+        model_id: MLXModelId,
+        messages: Sequence[Message],
+        tools: Sequence[Tool] | Toolkit | None = None,
+        format: type[FormattableT] | Format[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, tools, format, params
+        )
+
+        return Response(
+            raw=response,
+            provider_id="mlx",
+            model_id=model_id,
+            provider_model_name=model_id,
+            params=params,
+            tools=tools,
+            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],
+        tools: Sequence[Tool | ContextTool[DepsT]]
+        | ContextToolkit[DepsT]
+        | None = None,
+        format: type[FormattableT] | Format[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, tools, format, params
+        )
+
+        return ContextResponse(
+            raw=response,
+            provider_id="mlx",
+            model_id=model_id,
+            provider_model_name=model_id,
+            params=params,
+            tools=tools,
+            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],
+        tools: Sequence[AsyncTool] | AsyncToolkit | None = None,
+        format: type[FormattableT] | Format[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, tools, format, params)
+
+        return AsyncResponse(
+            raw=response,
+            provider_id="mlx",
+            model_id=model_id,
+            provider_model_name=model_id,
+            params=params,
+            tools=tools,
+            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],
+        tools: Sequence[AsyncTool | AsyncContextTool[DepsT]]
+        | AsyncContextToolkit[DepsT]
+        | None = None,
+        format: type[FormattableT] | Format[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, tools, format, params)
+
+        return AsyncContextResponse(
+            raw=response,
+            provider_id="mlx",
+            model_id=model_id,
+            provider_model_name=model_id,
+            params=params,
+            tools=tools,
+            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],
+        tools: Sequence[Tool] | Toolkit | None = None,
+        format: type[FormattableT] | Format[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, tools, format, params
+        )
+
+        return StreamResponse(
+            provider_id="mlx",
+            model_id=model_id,
+            provider_model_name=model_id,
+            params=params,
+            tools=tools,
+            input_messages=input_messages,
+            chunk_iterator=chunk_iterator,
+            format=format,
+        )
+
+    def _context_stream(
+        self,
+        *,
+        ctx: Context[DepsT],
+        model_id: MLXModelId,
+        messages: Sequence[Message],
+        tools: Sequence[Tool | ContextTool[DepsT]]
+        | ContextToolkit[DepsT]
+        | None = None,
+        format: type[FormattableT] | Format[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, tools, format, params
+        )
+
+        return ContextStreamResponse(
+            provider_id="mlx",
+            model_id=model_id,
+            provider_model_name=model_id,
+            params=params,
+            tools=tools,
+            input_messages=input_messages,
+            chunk_iterator=chunk_iterator,
+            format=format,
+        )
+
+    async def _stream_async(
+        self,
+        *,
+        model_id: MLXModelId,
+        messages: Sequence[Message],
+        tools: Sequence[AsyncTool] | AsyncToolkit | None = None,
+        format: type[FormattableT] | Format[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, tools, format, params
+        )
+
+        return AsyncStreamResponse(
+            provider_id="mlx",
+            model_id=model_id,
+            provider_model_name=model_id,
+            params=params,
+            tools=tools,
+            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],
+        tools: Sequence[AsyncTool | AsyncContextTool[DepsT]]
+        | AsyncContextToolkit[DepsT]
+        | None = None,
+        format: type[FormattableT] | Format[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, tools, format, params
+        )
+
+        return AsyncContextStreamResponse(
+            provider_id="mlx",
+            model_id=model_id,
+            provider_model_name=model_id,
+            params=params,
+            tools=tools,
+            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,19 @@
+"""Ollama provider implementation."""
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .provider import OllamaProvider
+else:
+    try:
+        from .provider import OllamaProvider
+    except ImportError:  # pragma: no cover
+        from .._missing_import_stubs import (
+            create_provider_stub,
+        )
+
+        OllamaProvider = create_provider_stub("openai", "OllamaProvider")
+
+__all__ = [
+    "OllamaProvider",
+]

mirascope/llm/providers/ollama/provider.py

@@ -0,0 +1,71 @@
+"""Ollama provider implementation."""
+
+import os
+from typing import ClassVar
+
+from openai import AsyncOpenAI, OpenAI
+
+from ..openai.completions.base_provider import BaseOpenAICompletionsProvider
+
+
+class OllamaProvider(BaseOpenAICompletionsProvider):
+    """Provider for Ollama's OpenAI-compatible API.
+
+    Inherits from BaseOpenAICompletionsProvider with Ollama-specific configuration:
+    - Uses Ollama's local API endpoint (default: http://localhost:11434/v1/)
+    - API key is not required (Ollama ignores API keys)
+    - Supports OLLAMA_BASE_URL environment variable
+
+    Usage:
+        Register the provider with model ID prefixes you want to use:
+
+        ```python
+        import llm
+
+        # Register for ollama models
+        llm.register_provider("ollama", "ollama/")
+
+        # Now you can use ollama models directly
+        @llm.call("ollama/llama2")
+        def my_prompt():
+            return [llm.messages.user("Hello!")]
+        ```
+    """
+
+    id: ClassVar[str] = "ollama"
+    default_scope: ClassVar[str | list[str]] = "ollama/"
+    default_base_url: ClassVar[str | None] = "http://localhost:11434/v1/"
+    api_key_env_var: ClassVar[str] = "OLLAMA_API_KEY"
+    api_key_required: ClassVar[bool] = False
+    provider_name: ClassVar[str | None] = "Ollama"
+
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+        base_url: str | None = None,
+    ) -> None:
+        """Initialize the Ollama provider.
+
+        Args:
+            api_key: API key (optional). Defaults to OLLAMA_API_KEY env var or 'ollama'.
+            base_url: Custom base URL. Defaults to OLLAMA_BASE_URL env var
+                or http://localhost:11434/v1/.
+        """
+        resolved_api_key = api_key or os.environ.get(self.api_key_env_var) or "ollama"
+        resolved_base_url = (
+            base_url or os.environ.get("OLLAMA_BASE_URL") or self.default_base_url
+        )
+
+        self.client = OpenAI(
+            api_key=resolved_api_key,
+            base_url=resolved_base_url,
+        )
+        self.async_client = AsyncOpenAI(
+            api_key=resolved_api_key,
+            base_url=resolved_base_url,
+        )
+
+    def _model_name(self, model_id: str) -> str:
+        """Strip 'ollama/' prefix from model ID for Ollama API."""
+        return model_id.removeprefix("ollama/")

mirascope/llm/providers/openai/__init__.py

@@ -0,0 +1,6 @@
+"""OpenAI client implementation."""
+
+from .model_id import OpenAIModelId
+from .provider import OpenAIProvider
+
+__all__ = ["OpenAIModelId", "OpenAIProvider"]

mirascope/llm/providers/openai/completions/__init__.py

@@ -0,0 +1,25 @@
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .base_provider import BaseOpenAICompletionsProvider
+    from .provider import OpenAICompletionsProvider
+else:
+    try:
+        from .base_provider import BaseOpenAICompletionsProvider
+        from .provider import OpenAICompletionsProvider
+    except ImportError:  # pragma: no cover
+        from ..._missing_import_stubs import (
+            create_provider_stub,
+        )
+
+        BaseOpenAICompletionsProvider = create_provider_stub(
+            "openai", "BaseOpenAICompletionsProvider"
+        )
+        OpenAICompletionsProvider = create_provider_stub(
+            "openai", "OpenAICompletionsProvider"
+        )
+
+__all__ = [
+    "BaseOpenAICompletionsProvider",
+    "OpenAICompletionsProvider",
+]

mirascope/llm/{clients → providers}/openai/completions/_utils/__init__.py

@@ -2,6 +2,7 @@ from .decode import (
     decode_async_stream,
     decode_response,
     decode_stream,
+    model_name,
 )
 from .encode import encode_request
 
@@ -10,4 +11,5 @@ __all__ = [
     "decode_response",
     "decode_stream",
     "encode_request",
+    "model_name",
 ]

mirascope/llm/{clients → providers}/openai/completions/_utils/decode.py

@@ -4,6 +4,7 @@ from typing import Literal
 
 from openai import AsyncStream, Stream
 from openai.types import chat as openai_types
+from openai.types.completion_usage import CompletionUsage
 
 from .....content import (
     AssistantContentPart,
@@ -23,8 +24,10 @@ from .....responses import (
     FinishReason,
     FinishReasonChunk,
     RawStreamEventChunk,
+    Usage,
+    UsageDeltaChunk,
 )
-from ..model_ids import OpenAICompletionsModelId
+from ...model_id import OpenAIModelId, model_name
 
 OPENAI_FINISH_REASON_MAP = {
     "length": FinishReason.MAX_TOKENS,
@@ -32,11 +35,40 @@ OPENAI_FINISH_REASON_MAP = {
 }
 
 
+def _decode_usage(
+    usage: CompletionUsage | None,
+) -> Usage | None:
+    """Convert OpenAI CompletionUsage to Mirascope Usage."""
+    if usage is None:  # pragma: no cover
+        return None
+
+    return Usage(
+        input_tokens=usage.prompt_tokens,
+        output_tokens=usage.completion_tokens,
+        cache_read_tokens=(
+            usage.prompt_tokens_details.cached_tokens
+            if usage.prompt_tokens_details
+            else None
+        )
+        or 0,
+        cache_write_tokens=0,
+        reasoning_tokens=(
+            usage.completion_tokens_details.reasoning_tokens
+            if usage.completion_tokens_details
+            else None
+        )
+        or 0,
+        raw=usage,
+    )
+
+
 def decode_response(
     response: openai_types.ChatCompletion,
-    model_id: OpenAICompletionsModelId,
-) -> tuple[AssistantMessage, FinishReason | None]:
-    """Convert OpenAI ChatCompletion to mirascope AssistantMessage."""
+    model_id: OpenAIModelId,
+    provider_id: str,
+    provider_model_name: str | None = None,
+) -> tuple[AssistantMessage, FinishReason | None, Usage | None]:
+    """Convert OpenAI ChatCompletion to mirascope AssistantMessage and usage."""
     choice = response.choices[0]
     message = choice.message
     refused = False
@@ -69,12 +101,14 @@ def decode_response(
 
     assistant_message = AssistantMessage(
         content=parts,
-        provider="openai:completions",
+        provider_id=provider_id,
         model_id=model_id,
+        provider_model_name=provider_model_name or model_name(model_id, "completions"),
         raw_message=message.model_dump(exclude_none=True),
     )
 
-    return assistant_message, finish_reason
+    usage = _decode_usage(response.usage)
+    return assistant_message, finish_reason, usage
 
 
 class _OpenAIChunkProcessor:
@@ -89,6 +123,26 @@ class _OpenAIChunkProcessor:
         """Process a single OpenAI chunk and yield the appropriate content chunks."""
         yield RawStreamEventChunk(raw_stream_event=chunk)
 
+        if chunk.usage:
+            usage = chunk.usage
+            yield UsageDeltaChunk(
+                input_tokens=usage.prompt_tokens,
+                output_tokens=usage.completion_tokens,
+                cache_read_tokens=(
+                    usage.prompt_tokens_details.cached_tokens
+                    if usage.prompt_tokens_details
+                    else None
+                )
+                or 0,
+                cache_write_tokens=0,
+                reasoning_tokens=(
+                    usage.completion_tokens_details.reasoning_tokens
+                    if usage.completion_tokens_details
+                    else None
+                )
+                or 0,
+            )
+
         choice = chunk.choices[0] if chunk.choices else None
         if not choice:
             return  # pragma: no cover