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

mirascope/llm/providers/google/__init__.py

@@ -0,0 +1,21 @@
+"""Google client implementation."""
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .model_id import GoogleModelId
+    from .provider import GoogleProvider
+else:
+    try:
+        from .model_id import GoogleModelId
+        from .provider import GoogleProvider
+    except ImportError:  # pragma: no cover
+        from .._missing_import_stubs import (
+            create_import_error_stub,
+            create_provider_stub,
+        )
+
+        GoogleProvider = create_provider_stub("google", "GoogleProvider")
+        GoogleModelId = str
+
+__all__ = ["GoogleModelId", "GoogleProvider"]

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

@@ -30,7 +30,7 @@ from ....responses import (
     RawMessageChunk,
     RawStreamEventChunk,
 )
-from ..model_ids import GoogleModelId
+from ..model_id import GoogleModelId, model_name
 from .encode import UNKNOWN_TOOL_ID
 
 GOOGLE_FINISH_REASON_MAP = {
@@ -88,7 +88,7 @@ def _decode_candidate_content(
     candidate: genai_types.Candidate,
 ) -> Sequence[AssistantContentPart]:
     """Returns a sequence of `AssistantContentPart` decoded from a google `Candidate`"""
-    content_parts = []
+    content_parts: list[AssistantContentPart] = []
     if candidate.content and candidate.content.parts:
         for part in candidate.content.parts:
             decoded_part = _decode_content_part(part)
@@ -98,7 +98,8 @@ def _decode_candidate_content(
 
 
 def decode_response(
-    response: genai_types.GenerateContentResponse, model_id: GoogleModelId
+    response: genai_types.GenerateContentResponse,
+    model_id: GoogleModelId,
 ) -> tuple[AssistantMessage, FinishReason | None]:
     """Returns an `AssistantMessage` and `FinishReason` extracted from a `GenerateContentResponse`"""
     content: Sequence[AssistantContentPart] = []
@@ -115,8 +116,9 @@ def decode_response(
 
     assistant_message = AssistantMessage(
         content=content,
-        provider="google",
+        provider_id="google",
         model_id=model_id,
+        provider_model_name=model_name(model_id),
         raw_message=candidate_content.model_dump(),
     )
 

mirascope/llm/{clients → providers}/google/_utils/encode.py

@@ -4,7 +4,8 @@ import base64
 import json
 from collections.abc import Sequence
 from functools import lru_cache
-from typing import Any, cast
+from typing import Any, TypedDict, cast
+from typing_extensions import Required
 
 from google.genai import types as genai_types
 
@@ -17,15 +18,19 @@ from ....formatting import (
     resolve_format,
 )
 from ....messages import AssistantMessage, Message, UserMessage
-from ....tools import FORMAT_TOOL_NAME, BaseToolkit, ToolSchema
-from ...base import BaseKwargs, Params, _utils as _base_utils
-from ..model_ids import GoogleModelId
+from ....tools import FORMAT_TOOL_NAME, AnyToolSchema, BaseToolkit
+from ...base import Params, _utils as _base_utils
+from ..model_id import GoogleModelId, model_name
 
 UNKNOWN_TOOL_ID = "google_unknown_tool_id"
 
 
-class GoogleKwargs(BaseKwargs, genai_types.GenerateContentConfigDict):
-    """Google's `GenerateContentConfigDict` typed dict, subclassing BaseKwargs for type safety."""
+class GoogleKwargs(TypedDict, total=False):
+    """Kwargs for Google's generate_content method."""
+
+    model: Required[str]
+    contents: Required[genai_types.ContentListUnionDict]
+    config: genai_types.GenerateContentConfigDict
 
 
 def _resolve_refs(
@@ -52,7 +57,7 @@ def _encode_content(
     content: Sequence[ContentPart], encode_thoughts: bool
 ) -> list[genai_types.PartDict]:
     """Returns a list of google `PartDicts` converted from a sequence of Mirascope `ContentPart`s"""
-    result = []
+    result: list[genai_types.PartDict] = []
 
     for part in content:
         if part.type == "text":
@@ -121,7 +126,7 @@ def _encode_message(
     """Returns a Google `ContentDict` converted from a Mirascope `Message`"""
     if (
         message.role == "assistant"
-        and message.provider == "google"
+        and message.provider_id == "google"
         and message.model_id == model_id
         and message.raw_message
         and not encode_thoughts
@@ -144,7 +149,7 @@ def _encode_messages(
 
 @lru_cache(maxsize=128)
 def _convert_tool_to_function_declaration(
-    tool: ToolSchema,
+    tool: AnyToolSchema,
 ) -> genai_types.FunctionDeclarationDict:
     """Convert a single Mirascope tool to Google FunctionDeclaration format with caching."""
     schema_dict = tool.parameters.model_dump(by_alias=True, exclude_none=True)
@@ -170,21 +175,21 @@ def encode_request(
     *,
     model_id: GoogleModelId,
     messages: Sequence[Message],
-    tools: Sequence[ToolSchema] | BaseToolkit | None,
+    tools: Sequence[AnyToolSchema] | BaseToolkit[AnyToolSchema] | None,
     format: type[FormattableT] | Format[FormattableT] | None,
     params: Params,
-) -> tuple[
-    Sequence[Message],
-    Format[FormattableT] | None,
-    genai_types.ContentListUnionDict,
-    GoogleKwargs,
-]:
+) -> tuple[Sequence[Message], Format[FormattableT] | None, GoogleKwargs]:
     """Prepares a request for the genai `Client.models.generate_content` method."""
-    google_config: GoogleKwargs = GoogleKwargs()
+    if not model_id.startswith("google/"):  # pragma: no cover
+        raise ValueError(f"Model ID must start with 'google/' prefix, got: {model_id}")
+
+    google_config: genai_types.GenerateContentConfigDict = (
+        genai_types.GenerateContentConfigDict()
+    )
     encode_thoughts = False
 
     with _base_utils.ensure_all_params_accessed(
-        params=params, provider="google"
+        params=params, provider_id="google"
     ) as param_accessor:
         if param_accessor.temperature is not None:
             google_config["temperature"] = param_accessor.temperature
@@ -224,7 +229,7 @@ def encode_request(
         if format.mode in ("strict", "json") and tools:
             raise FeatureNotSupportedError(
                 feature=f"formatting_mode:{format.mode} with tools",
-                provider="google",
+                provider_id="google",
             )
 
         if format.mode == "strict":
@@ -271,9 +276,10 @@ def encode_request(
     if system_message_content:
         google_config["system_instruction"] = system_message_content
 
-    return (
-        messages,
-        format,
-        _encode_messages(remaining_messages, model_id, encode_thoughts),
-        google_config,
+    kwargs = GoogleKwargs(
+        model=model_name(model_id),
+        contents=_encode_messages(remaining_messages, model_id, encode_thoughts),
+        config=google_config,
     )
+
+    return messages, format, kwargs

mirascope/llm/providers/google/model_id.py

@@ -0,0 +1,28 @@
+"""Google registered LLM models."""
+
+from typing import Literal, TypeAlias
+
+GoogleModelId: TypeAlias = (
+    Literal[
+        "google/gemini-3-pro-preview",
+        "google/gemini-2.5-pro",
+        "google/gemini-2.5-flash",
+        "google/gemini-2.5-flash-lite",
+        "google/gemini-2.0-flash",
+        "google/gemini-2.0-flash-lite",
+    ]
+    | str
+)
+"""The Google model ids registered with Mirascope."""
+
+
+def model_name(model_id: GoogleModelId) -> str:
+    """Extract the google model name from a full model ID.
+
+    Args:
+        model_id: Full model ID (e.g. "google/gemini-2.5-flash")
+
+    Returns:
+        Provider-specific model ID (e.g. "gemini-2.5-flash")
+    """
+    return model_id.removeprefix("google/")

mirascope/llm/providers/google/provider.py

@@ -0,0 +1,438 @@
+"""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 = _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,
+            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 = _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,
+            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 = _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,
+            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 = _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,
+            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,48 @@
+from functools import lru_cache
+
+from .anthropic import AnthropicProvider
+from .base import Provider
+from .google import GoogleProvider
+from .mlx import MLXProvider
+from .openai import OpenAIProvider
+from .openai.completions.provider import OpenAICompletionsProvider
+from .openai.responses.provider import OpenAIResponsesProvider
+from .provider_id import ProviderId
+
+
+@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 "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 "mlx":  # pragma: no cover (MLX is only available on macOS)
+            return MLXProvider()
+        case _:  # pragma: no cover
+            raise ValueError(f"Unknown provider: '{provider_id}'")
+
+
+load = load_provider
+"""Convenient alias as `llm.providers.load`"""