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

mirascope/llm/providers/google/provider.py

@@ -0,0 +1,456 @@
+"""Google provider implementation."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
+from typing_extensions import Unpack
+
+from google.genai import Client
+from google.genai.types import HttpOptions
+
+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 .model_id import GoogleModelId, model_name
+
+if TYPE_CHECKING:
+    from ...models import Params
+
+
+class GoogleProvider(BaseProvider[Client]):
+    """The client for the Google LLM model."""
+
+    id = "google"
+    default_scope = "google/"
+    error_map = _utils.GOOGLE_ERROR_MAP
+
+    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,
+                headers={"Authorization": f"Bearer {api_key}"},
+            )
+
+        self.client = Client(api_key=api_key, http_options=http_options)
+
+    def get_error_status(self, e: Exception) -> int | None:
+        """Extract HTTP status code from Google exception."""
+        return getattr(e, "code", None)
+
+    def _call(
+        self,
+        *,
+        model_id: GoogleModelId,
+        messages: Sequence[Message],
+        toolkit: Toolkit,
+        format: FormatSpec[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=toolkit,
+            format=format,
+            params=params,
+        )
+        google_response = self.client.models.generate_content(**kwargs)
+
+        include_thoughts = _utils.get_include_thoughts(params)
+        assistant_message, finish_reason, usage = _utils.decode_response(
+            google_response, model_id, include_thoughts=include_thoughts
+        )
+
+        return Response(
+            raw=google_response,
+            provider_id="google",
+            model_id=model_id,
+            provider_model_name=model_name(model_id),
+            params=params,
+            tools=toolkit,
+            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],
+        toolkit: ContextToolkit[DepsT],
+        format: FormatSpec[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=toolkit,
+            format=format,
+            params=params,
+        )
+        google_response = self.client.models.generate_content(**kwargs)
+
+        include_thoughts = _utils.get_include_thoughts(params)
+        assistant_message, finish_reason, usage = _utils.decode_response(
+            google_response, model_id, include_thoughts=include_thoughts
+        )
+
+        return ContextResponse(
+            raw=google_response,
+            provider_id="google",
+            model_id=model_id,
+            provider_model_name=model_name(model_id),
+            params=params,
+            tools=toolkit,
+            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],
+        toolkit: AsyncToolkit,
+        format: FormatSpec[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=toolkit,
+            format=format,
+            params=params,
+        )
+        google_response = await self.client.aio.models.generate_content(**kwargs)
+
+        include_thoughts = _utils.get_include_thoughts(params)
+        assistant_message, finish_reason, usage = _utils.decode_response(
+            google_response, model_id, include_thoughts=include_thoughts
+        )
+
+        return AsyncResponse(
+            raw=google_response,
+            provider_id="google",
+            model_id=model_id,
+            provider_model_name=model_name(model_id),
+            params=params,
+            tools=toolkit,
+            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],
+        toolkit: AsyncContextToolkit[DepsT],
+        format: FormatSpec[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=toolkit,
+            format=format,
+            params=params,
+        )
+        google_response = await self.client.aio.models.generate_content(**kwargs)
+
+        include_thoughts = _utils.get_include_thoughts(params)
+        assistant_message, finish_reason, usage = _utils.decode_response(
+            google_response, model_id, include_thoughts=include_thoughts
+        )
+
+        return AsyncContextResponse(
+            raw=google_response,
+            provider_id="google",
+            model_id=model_id,
+            provider_model_name=model_name(model_id),
+            params=params,
+            tools=toolkit,
+            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],
+        toolkit: Toolkit,
+        format: FormatSpec[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=toolkit,
+            format=format,
+            params=params,
+        )
+
+        google_stream = self.client.models.generate_content_stream(**kwargs)
+
+        include_thoughts = _utils.get_include_thoughts(params)
+        chunk_iterator = _utils.decode_stream(
+            google_stream, include_thoughts=include_thoughts
+        )
+
+        return StreamResponse(
+            provider_id="google",
+            model_id=model_id,
+            provider_model_name=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: GoogleModelId,
+        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 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=toolkit,
+            format=format,
+            params=params,
+        )
+
+        google_stream = self.client.models.generate_content_stream(**kwargs)
+
+        include_thoughts = _utils.get_include_thoughts(params)
+        chunk_iterator = _utils.decode_stream(
+            google_stream, include_thoughts=include_thoughts
+        )
+
+        return ContextStreamResponse(
+            provider_id="google",
+            model_id=model_id,
+            provider_model_name=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: GoogleModelId,
+        messages: Sequence[Message],
+        toolkit: AsyncToolkit,
+        format: FormatSpec[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=toolkit,
+            format=format,
+            params=params,
+        )
+
+        google_stream = await self.client.aio.models.generate_content_stream(**kwargs)
+
+        include_thoughts = _utils.get_include_thoughts(params)
+        chunk_iterator = _utils.decode_async_stream(
+            google_stream, include_thoughts=include_thoughts
+        )
+
+        return AsyncStreamResponse(
+            provider_id="google",
+            model_id=model_id,
+            provider_model_name=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: GoogleModelId,
+        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 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=toolkit,
+            format=format,
+            params=params,
+        )
+
+        google_stream = await self.client.aio.models.generate_content_stream(**kwargs)
+
+        include_thoughts = _utils.get_include_thoughts(params)
+        chunk_iterator = _utils.decode_async_stream(
+            google_stream, include_thoughts=include_thoughts
+        )
+
+        return AsyncContextStreamResponse(
+            provider_id="google",
+            model_id=model_id,
+            provider_model_name=model_name(model_id),
+            params=params,
+            tools=toolkit,
+            input_messages=input_messages,
+            chunk_iterator=chunk_iterator,
+            format=format,
+        )

mirascope/llm/providers/mirascope/__init__.py

@@ -0,0 +1,5 @@
+"""Mirascope Router provider for routing LLM requests."""
+
+from .provider import MirascopeProvider
+
+__all__ = ["MirascopeProvider"]

mirascope/llm/providers/mirascope/_utils.py

@@ -0,0 +1,73 @@
+"""Utility functions for Mirascope Router provider."""
+
+import os
+from typing import cast
+
+from ..base import Provider
+from ..provider_id import ProviderId
+
+
+def extract_model_scope(model_id: str) -> str | None:
+    """Extract model scope from model ID.
+
+    Args:
+        model_id: Model identifier in the format "scope/model-name"
+                  e.g., "openai/gpt-4", "anthropic/claude-3", "google/gemini-pro"
+
+    Returns:
+        The model scope (e.g., "openai", "anthropic", "google") or None if invalid format.
+    """
+    if "/" not in model_id:
+        return None
+    return model_id.split("/", 1)[0]
+
+
+def get_default_router_base_url() -> str:
+    """Get the default router base URL from environment or use default.
+
+    Returns:
+        The router base URL (without trailing provider path).
+    """
+    return os.environ.get(
+        "MIRASCOPE_ROUTER_BASE_URL", "https://mirascope.com/router/v2"
+    )
+
+
+def create_underlying_provider(
+    model_scope: str, api_key: str, router_base_url: str
+) -> Provider:
+    """Create and cache an underlying provider instance using provider_singleton.
+
+    This function constructs the appropriate router URL for the provider and
+    delegates to provider_singleton for caching and instantiation.
+
+    Args:
+        model_scope: The model scope (e.g., "openai", "anthropic", "google")
+        api_key: The API key to use for authentication
+        router_base_url: The base router URL (e.g., "http://mirascope.com/router/v2")
+
+    Returns:
+        A cached provider instance configured for the Mirascope Router.
+
+    Raises:
+        ValueError: If the provider is unsupported.
+    """
+    if model_scope not in ["anthropic", "google", "openai"]:
+        raise ValueError(
+            f"Unsupported provider: {model_scope}. "
+            f"Mirascope Router currently supports: anthropic, google, openai"
+        )
+
+    base_url = f"{router_base_url}/{model_scope}"
+    if model_scope == "openai":  # OpenAI expects /v1, which their SDK doesn't add
+        base_url = f"{base_url}/v1"
+
+    # Lazy import to avoid circular dependencies
+    from ..provider_registry import provider_singleton
+
+    # Use provider_singleton which provides caching
+    return provider_singleton(
+        cast(ProviderId, model_scope),
+        api_key=api_key,
+        base_url=base_url,
+    )