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

mirascope/llm/providers/google/provider.py

@@ -0,0 +1,442 @@
+"""Google provider implementation."""
+
+from collections.abc import Sequence
+from typing_extensions import Unpack
+
+from google.genai import Client
+from google.genai.types import HttpOptions
+
+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 .model_id import GoogleModelId, model_name
+
+
+class GoogleProvider(BaseProvider[Client]):
+    """The client for the Google LLM model."""
+
+    id = "google"
+    default_scope = "google/"
+
+    def __init__(
+        self, *, api_key: str | None = None, base_url: str | None = None
+    ) -> None:
+        """Initialize the Google client."""
+        http_options = None
+        if base_url:
+            http_options = HttpOptions(base_url=base_url)
+
+        self.client = Client(api_key=api_key, http_options=http_options)
+
+    def _call(
+        self,
+        *,
+        model_id: GoogleModelId,
+        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` by synchronously calling the Google GenAI API.
+
+        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.
+        """
+        input_messages, format, kwargs = _utils.encode_request(
+            model_id=model_id,
+            messages=messages,
+            tools=tools,
+            format=format,
+            params=params,
+        )
+
+        google_response = self.client.models.generate_content(**kwargs)
+
+        assistant_message, finish_reason, usage = _utils.decode_response(
+            google_response, model_id
+        )
+
+        return Response(
+            raw=google_response,
+            provider_id="google",
+            model_id=model_id,
+            provider_model_name=model_name(model_id),
+            params=params,
+            tools=tools,
+            input_messages=input_messages,
+            assistant_message=assistant_message,
+            finish_reason=finish_reason,
+            usage=usage,
+            format=format,
+        )
+
+    def _context_call(
+        self,
+        *,
+        ctx: Context[DepsT],
+        model_id: GoogleModelId,
+        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` by synchronously calling the Google GenAI API.
+
+        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.
+        """
+        input_messages, format, kwargs = _utils.encode_request(
+            model_id=model_id,
+            messages=messages,
+            tools=tools,
+            format=format,
+            params=params,
+        )
+
+        google_response = self.client.models.generate_content(**kwargs)
+
+        assistant_message, finish_reason, usage = _utils.decode_response(
+            google_response, model_id
+        )
+
+        return ContextResponse(
+            raw=google_response,
+            provider_id="google",
+            model_id=model_id,
+            provider_model_name=model_name(model_id),
+            params=params,
+            tools=tools,
+            input_messages=input_messages,
+            assistant_message=assistant_message,
+            finish_reason=finish_reason,
+            usage=usage,
+            format=format,
+        )
+
+    async def _call_async(
+        self,
+        *,
+        model_id: GoogleModelId,
+        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` by asynchronously calling the Google GenAI API.
+
+        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.
+        """
+        input_messages, format, kwargs = _utils.encode_request(
+            model_id=model_id,
+            messages=messages,
+            tools=tools,
+            format=format,
+            params=params,
+        )
+
+        google_response = await self.client.aio.models.generate_content(**kwargs)
+
+        assistant_message, finish_reason, usage = _utils.decode_response(
+            google_response, model_id
+        )
+
+        return AsyncResponse(
+            raw=google_response,
+            provider_id="google",
+            model_id=model_id,
+            provider_model_name=model_name(model_id),
+            params=params,
+            tools=tools,
+            input_messages=input_messages,
+            assistant_message=assistant_message,
+            finish_reason=finish_reason,
+            usage=usage,
+            format=format,
+        )
+
+    async def _context_call_async(
+        self,
+        *,
+        ctx: Context[DepsT],
+        model_id: GoogleModelId,
+        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.AsyncContextResponse` by asynchronously calling the Google GenAI API.
+
+        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.
+        """
+        input_messages, format, kwargs = _utils.encode_request(
+            model_id=model_id,
+            messages=messages,
+            tools=tools,
+            format=format,
+            params=params,
+        )
+
+        google_response = await self.client.aio.models.generate_content(**kwargs)
+
+        assistant_message, finish_reason, usage = _utils.decode_response(
+            google_response, model_id
+        )
+
+        return AsyncContextResponse(
+            raw=google_response,
+            provider_id="google",
+            model_id=model_id,
+            provider_model_name=model_name(model_id),
+            params=params,
+            tools=tools,
+            input_messages=input_messages,
+            assistant_message=assistant_message,
+            finish_reason=finish_reason,
+            usage=usage,
+            format=format,
+        )
+
+    def _stream(
+        self,
+        *,
+        model_id: GoogleModelId,
+        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 the Google GenAI API.
+
+        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.
+        """
+        input_messages, format, kwargs = _utils.encode_request(
+            model_id=model_id,
+            messages=messages,
+            tools=tools,
+            format=format,
+            params=params,
+        )
+
+        google_stream = self.client.models.generate_content_stream(**kwargs)
+
+        chunk_iterator = _utils.decode_stream(google_stream)
+
+        return StreamResponse(
+            provider_id="google",
+            model_id=model_id,
+            provider_model_name=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: GoogleModelId,
+        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 the Google GenAI API.
+
+        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.
+        """
+        input_messages, format, kwargs = _utils.encode_request(
+            model_id=model_id,
+            messages=messages,
+            tools=tools,
+            format=format,
+            params=params,
+        )
+
+        google_stream = self.client.models.generate_content_stream(**kwargs)
+
+        chunk_iterator = _utils.decode_stream(google_stream)
+
+        return ContextStreamResponse(
+            provider_id="google",
+            model_id=model_id,
+            provider_model_name=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: GoogleModelId,
+        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 the Google GenAI API.
+
+        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.
+        """
+        input_messages, format, kwargs = _utils.encode_request(
+            model_id=model_id,
+            messages=messages,
+            tools=tools,
+            format=format,
+            params=params,
+        )
+
+        google_stream = await self.client.aio.models.generate_content_stream(**kwargs)
+
+        chunk_iterator = _utils.decode_async_stream(google_stream)
+
+        return AsyncStreamResponse(
+            provider_id="google",
+            model_id=model_id,
+            provider_model_name=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: GoogleModelId,
+        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 the Google GenAI API.
+
+        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.
+        """
+        input_messages, format, kwargs = _utils.encode_request(
+            model_id=model_id,
+            messages=messages,
+            tools=tools,
+            format=format,
+            params=params,
+        )
+
+        google_stream = await self.client.aio.models.generate_content_stream(**kwargs)
+
+        chunk_iterator = _utils.decode_async_stream(google_stream)
+
+        return AsyncContextStreamResponse(
+            provider_id="google",
+            model_id=model_id,
+            provider_model_name=model_name(model_id),
+            params=params,
+            tools=tools,
+            input_messages=input_messages,
+            chunk_iterator=chunk_iterator,
+            format=format,
+        )

mirascope/llm/providers/load_provider.py

@@ -0,0 +1,54 @@
+from functools import lru_cache
+
+from .anthropic import AnthropicProvider
+from .base import Provider
+from .google import GoogleProvider
+from .mlx import MLXProvider
+from .ollama import OllamaProvider
+from .openai import OpenAIProvider
+from .openai.completions.provider import OpenAICompletionsProvider
+from .openai.responses.provider import OpenAIResponsesProvider
+from .provider_id import ProviderId
+from .together import TogetherProvider
+
+
+@lru_cache(maxsize=256)
+def load_provider(
+    provider_id: ProviderId, *, api_key: str | None = None, base_url: str | None = None
+) -> Provider:
+    """Create a cached provider instance for the specified provider id.
+
+    Args:
+        provider_id: The provider name ("openai", "anthropic", or "google").
+        api_key: API key for authentication. If None, uses provider-specific env var.
+        base_url: Base URL for the API. If None, uses provider-specific env var.
+
+    Returns:
+        A cached provider instance for the specified provider with the given parameters.
+
+    Raises:
+        ValueError: If the provider_id is not supported.
+    """
+    match provider_id:
+        case "anthropic":
+            return AnthropicProvider(api_key=api_key, base_url=base_url)
+        case "google":
+            return GoogleProvider(api_key=api_key, base_url=base_url)
+        case "mlx":  # pragma: no cover (MLX is only available on macOS)
+            return MLXProvider()
+        case "ollama":
+            return OllamaProvider(api_key=api_key, base_url=base_url)
+        case "openai":
+            return OpenAIProvider(api_key=api_key, base_url=base_url)
+        case "openai:completions":
+            return OpenAICompletionsProvider(api_key=api_key, base_url=base_url)
+        case "openai:responses":
+            return OpenAIResponsesProvider(api_key=api_key, base_url=base_url)
+        case "together":
+            return TogetherProvider(api_key=api_key, base_url=base_url)
+        case _:  # pragma: no cover
+            raise ValueError(f"Unknown provider: '{provider_id}'")
+
+
+load = load_provider
+"""Convenient alias as `llm.providers.load`"""

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,129 @@
+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, Usage
+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
+
+
+def extract_usage(response: GenerationResponse | None) -> Usage | None:
+    """Extract usage information from an MLX generation response.
+
+    Args:
+        response: The MLX generation response to extract from.
+
+    Returns:
+        The Usage object with token counts, or None if not applicable.
+    """
+    if response is None:
+        return None
+
+    return Usage(
+        input_tokens=response.prompt_tokens,
+        output_tokens=response.generation_tokens,
+        cache_read_tokens=0,
+        cache_write_tokens=0,
+        reasoning_tokens=0,
+        raw=response,
+    )

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

@@ -0,0 +1,8 @@
+from .base import BaseEncoder, TokenIds
+from .transformers import TransformersEncoder
+
+__all__ = [
+    "BaseEncoder",
+    "TokenIds",
+    "TransformersEncoder",
+]