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

mirascope/llm/providers/base/base_provider.py

@@ -3,18 +3,22 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from collections.abc import Sequence
-from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeAlias, overload
+from collections.abc import Callable, Generator, Mapping, Sequence
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeAlias, cast, overload
 from typing_extensions import TypeVar, Unpack
 
 from ...context import Context, DepsT
+from ...exceptions import APIError, MirascopeLLMError
 from ...formatting import Format, FormattableT
 from ...messages import Message, UserContent, user
 from ...responses import (
+    AsyncChunkIterator,
     AsyncContextResponse,
     AsyncContextStreamResponse,
     AsyncResponse,
     AsyncStreamResponse,
+    ChunkIterator,
     ContextResponse,
     ContextStreamResponse,
     Response,
@@ -33,6 +37,7 @@ from ...tools import (
 from .params import Params
 
 if TYPE_CHECKING:
+    from ...exceptions import MirascopeLLMError
     from ..provider_id import ProviderId
 
 ProviderClientT = TypeVar("ProviderClientT")
@@ -40,6 +45,18 @@ ProviderClientT = TypeVar("ProviderClientT")
 Provider: TypeAlias = "BaseProvider[Any]"
 """Type alias for `BaseProvider` with any client type."""
 
+ProviderErrorMap: TypeAlias = Mapping[
+    type[Exception],
+    "type[MirascopeLLMError] | Callable[[Exception], type[MirascopeLLMError]]",
+]
+"""Mapping from provider SDK exceptions to Mirascope error types.
+
+Keys are provider SDK exception types (e.g., OpenAIError, AnthropicError).
+Values can be:
+- Error type: Simple 1:1 mapping (e.g., RateLimitError)
+- Callable: Transform function returning error type based on exception details
+"""
+
 
 class BaseProvider(Generic[ProviderClientT], ABC):
     """Base abstract provider for LLM interactions.
@@ -59,8 +76,67 @@ class BaseProvider(Generic[ProviderClientT], ABC):
     - ["anthropic/", "openai/"] - Multiple scopes (e.g., for AWS Bedrock)
     """
 
+    error_map: ClassVar[ProviderErrorMap]
+    """Mapping from provider SDK exceptions to Mirascope error types.
+
+    Values can be:
+    - Error type: Simple 1:1 mapping (e.g., AnthropicRateLimitError -> RateLimitError)
+    - Callable: Transform function returning error type based on exception details
+                (e.g., lambda e: NotFoundError if e.code == "model_not_found" else BadRequestError)
+
+    The mapping is walked via the exception's MRO, allowing both specific error handling
+    and fallback to base SDK error types (e.g., AnthropicError -> APIError).
+    """
+
     client: ProviderClientT
 
+    @contextmanager
+    def _wrap_errors(self) -> Generator[None, None, None]:
+        """Wrap provider API calls and convert errors to Mirascope exceptions.
+
+        Walks the exception's MRO to find the first matching error type in the
+        provider's error_map, allowing both specific error handling and fallback
+        to base SDK error types (e.g., AnthropicError -> APIError).
+        """
+        try:
+            yield
+        except Exception as e:
+            # Walk MRO to find first matching error type in provider's error_map
+            for error_class in type(e).__mro__:
+                if error_class in self.error_map:
+                    error_type_or_fn = self.error_map[error_class]
+
+                    if isinstance(error_type_or_fn, type):
+                        error_type = cast(type[MirascopeLLMError], error_type_or_fn)
+                    else:
+                        error_type = error_type_or_fn(e)
+
+                    # Construct Mirascope error with metadata
+                    error: MirascopeLLMError = error_type(str(e))
+                    if isinstance(error, APIError):
+                        error.status_code = self.get_error_status(e)
+                    error.provider = self.id
+                    error.original_exception = e
+                    raise error from e
+
+            # Not in error_map - not a provider error, re-raise as-is
+            raise
+
+    def _wrap_iterator_errors(self, iterator: ChunkIterator) -> ChunkIterator:
+        """Wrap sync chunk iterator to handle errors during iteration."""
+        # TODO: Consider moving this logic into BaseSyncStreamResponse if appropriate.
+        with self._wrap_errors():
+            yield from iterator
+
+    async def _wrap_async_iterator_errors(
+        self, iterator: AsyncChunkIterator
+    ) -> AsyncChunkIterator:
+        """Wrap async chunk iterator to handle errors during iteration."""
+        # TODO: Consider moving this logic into BaseAsyncStreamResponse if appropriate.
+        with self._wrap_errors():
+            async for chunk in iterator:
+                yield chunk
+
     @overload
     def call(
         self,
@@ -121,13 +197,14 @@ class BaseProvider(Generic[ProviderClientT], ABC):
         Returns:
             An `llm.Response` object containing the LLM-generated content.
         """
-        return self._call(
-            model_id=model_id,
-            messages=messages,
-            tools=tools,
-            format=format,
-            **params,
-        )
+        with self._wrap_errors():
+            return self._call(
+                model_id=model_id,
+                messages=messages,
+                tools=tools,
+                format=format,
+                **params,
+            )
 
     @abstractmethod
     def _call(
@@ -215,14 +292,15 @@ class BaseProvider(Generic[ProviderClientT], ABC):
         Returns:
             An `llm.ContextResponse` object containing the LLM-generated content.
         """
-        return self._context_call(
-            ctx=ctx,
-            model_id=model_id,
-            messages=messages,
-            tools=tools,
-            format=format,
-            **params,
-        )
+        with self._wrap_errors():
+            return self._context_call(
+                ctx=ctx,
+                model_id=model_id,
+                messages=messages,
+                tools=tools,
+                format=format,
+                **params,
+            )
 
     @abstractmethod
     def _context_call(
@@ -300,13 +378,14 @@ class BaseProvider(Generic[ProviderClientT], ABC):
         Returns:
             An `llm.AsyncResponse` object containing the LLM-generated content.
         """
-        return await self._call_async(
-            model_id=model_id,
-            messages=messages,
-            tools=tools,
-            format=format,
-            **params,
-        )
+        with self._wrap_errors():
+            return await self._call_async(
+                model_id=model_id,
+                messages=messages,
+                tools=tools,
+                format=format,
+                **params,
+            )
 
     @abstractmethod
     async def _call_async(
@@ -394,14 +473,15 @@ class BaseProvider(Generic[ProviderClientT], ABC):
         Returns:
             An `llm.AsyncContextResponse` object containing the LLM-generated content.
         """
-        return await self._context_call_async(
-            ctx=ctx,
-            model_id=model_id,
-            messages=messages,
-            tools=tools,
-            format=format,
-            **params,
-        )
+        with self._wrap_errors():
+            return await self._context_call_async(
+                ctx=ctx,
+                model_id=model_id,
+                messages=messages,
+                tools=tools,
+                format=format,
+                **params,
+            )
 
     @abstractmethod
     async def _context_call_async(
@@ -479,13 +559,18 @@ class BaseProvider(Generic[ProviderClientT], ABC):
         Returns:
             An `llm.StreamResponse` object for iterating over the LLM-generated content.
         """
-        return self._stream(
-            model_id=model_id,
-            messages=messages,
-            tools=tools,
-            format=format,
-            **params,
+        with self._wrap_errors():
+            stream_response = self._stream(
+                model_id=model_id,
+                messages=messages,
+                tools=tools,
+                format=format,
+                **params,
+            )
+        stream_response._chunk_iterator = self._wrap_iterator_errors(  # pyright: ignore[reportPrivateUsage]
+            stream_response._chunk_iterator  # pyright: ignore[reportPrivateUsage]
         )
+        return stream_response
 
     @abstractmethod
     def _stream(
@@ -577,14 +662,19 @@ class BaseProvider(Generic[ProviderClientT], ABC):
         Returns:
             An `llm.ContextStreamResponse` object for iterating over the LLM-generated content.
         """
-        return self._context_stream(
-            ctx=ctx,
-            model_id=model_id,
-            messages=messages,
-            tools=tools,
-            format=format,
-            **params,
+        with self._wrap_errors():
+            stream_response = self._context_stream(
+                ctx=ctx,
+                model_id=model_id,
+                messages=messages,
+                tools=tools,
+                format=format,
+                **params,
+            )
+        stream_response._chunk_iterator = self._wrap_iterator_errors(  # pyright: ignore[reportPrivateUsage]
+            stream_response._chunk_iterator  # pyright: ignore[reportPrivateUsage]
         )
+        return stream_response
 
     @abstractmethod
     def _context_stream(
@@ -664,13 +754,18 @@ class BaseProvider(Generic[ProviderClientT], ABC):
         Returns:
             An `llm.AsyncStreamResponse` object for asynchronously iterating over the LLM-generated content.
         """
-        return await self._stream_async(
-            model_id=model_id,
-            messages=messages,
-            tools=tools,
-            format=format,
-            **params,
+        with self._wrap_errors():
+            stream_response = await self._stream_async(
+                model_id=model_id,
+                messages=messages,
+                tools=tools,
+                format=format,
+                **params,
+            )
+        stream_response._chunk_iterator = self._wrap_async_iterator_errors(  # pyright: ignore[reportPrivateUsage]
+            stream_response._chunk_iterator  # pyright: ignore[reportPrivateUsage]
         )
+        return stream_response
 
     @abstractmethod
     async def _stream_async(
@@ -764,14 +859,19 @@ class BaseProvider(Generic[ProviderClientT], ABC):
         Returns:
             An `llm.AsyncContextStreamResponse` object for asynchronously iterating over the LLM-generated content.
         """
-        return await self._context_stream_async(
-            ctx=ctx,
-            model_id=model_id,
-            messages=messages,
-            tools=tools,
-            format=format,
-            **params,
+        with self._wrap_errors():
+            stream_response = await self._context_stream_async(
+                ctx=ctx,
+                model_id=model_id,
+                messages=messages,
+                tools=tools,
+                format=format,
+                **params,
+            )
+        stream_response._chunk_iterator = self._wrap_async_iterator_errors(  # pyright: ignore[reportPrivateUsage]
+            stream_response._chunk_iterator  # pyright: ignore[reportPrivateUsage]
         )
+        return stream_response
 
     @abstractmethod
     async def _context_stream_async(
@@ -1383,3 +1483,18 @@ class BaseProvider(Generic[ProviderClientT], ABC):
             format=response.format,
             **params,
         )
+
+    @abstractmethod
+    def get_error_status(self, e: Exception) -> int | None:
+        """Extract HTTP status code from provider-specific exception.
+
+        Different SDKs store status codes differently (e.g., .status_code vs .code).
+        Each provider implements this to handle their SDK's convention.
+
+        Args:
+            e: The exception to extract status code from.
+
+        Returns:
+            The HTTP status code if available, None otherwise.
+        """
+        ...

mirascope/llm/providers/google/_utils/__init__.py

@@ -4,8 +4,10 @@ from .decode import (
     decode_stream,
 )
 from .encode import encode_request
+from .errors import GOOGLE_ERROR_MAP
 
 __all__ = [
+    "GOOGLE_ERROR_MAP",
     "decode_async_stream",
     "decode_response",
     "decode_stream",

mirascope/llm/providers/google/_utils/decode.py

@@ -29,6 +29,8 @@ from ....responses import (
     FinishReasonChunk,
     RawMessageChunk,
     RawStreamEventChunk,
+    Usage,
+    UsageDeltaChunk,
 )
 from ..model_id import GoogleModelId, model_name
 from .encode import UNKNOWN_TOOL_ID
@@ -43,6 +45,30 @@ GOOGLE_FINISH_REASON_MAP = {
 }
 
 
+def _decode_usage(
+    usage: genai_types.GenerateContentResponseUsageMetadata | None,
+) -> Usage | None:
+    """Convert Google UsageMetadata to Mirascope Usage."""
+    if (
+        usage is None
+        or usage.prompt_token_count is None
+        or usage.candidates_token_count is None
+    ):  # pragma: no cover
+        return None
+
+    reasoning_tokens = usage.thoughts_token_count or 0
+    output_tokens = usage.candidates_token_count + reasoning_tokens
+
+    return Usage(
+        input_tokens=usage.prompt_token_count,
+        output_tokens=output_tokens,
+        cache_read_tokens=usage.cached_content_token_count or 0,
+        cache_write_tokens=0,
+        reasoning_tokens=usage.thoughts_token_count or 0,
+        raw=usage,
+    )
+
+
 def _decode_content_part(part: genai_types.Part) -> AssistantContentPart | None:
     """Returns an `AssistantContentPart` (or `None`) decoded from a google `Part`"""
     if part.thought and part.text:
@@ -100,8 +126,8 @@ def _decode_candidate_content(
 def decode_response(
     response: genai_types.GenerateContentResponse,
     model_id: GoogleModelId,
-) -> tuple[AssistantMessage, FinishReason | None]:
-    """Returns an `AssistantMessage` and `FinishReason` extracted from a `GenerateContentResponse`"""
+) -> tuple[AssistantMessage, FinishReason | None, Usage | None]:
+    """Returns an `AssistantMessage`, `FinishReason`, and `Usage` extracted from a `GenerateContentResponse`"""
     content: Sequence[AssistantContentPart] = []
     candidate_content: genai_types.Content | None = None
     finish_reason: FinishReason | None = None
@@ -122,7 +148,8 @@ def decode_response(
         raw_message=candidate_content.model_dump(),
     )
 
-    return assistant_message, finish_reason
+    usage = _decode_usage(response.usage_metadata)
+    return assistant_message, finish_reason, usage
 
 
 class _GoogleChunkProcessor:
@@ -132,6 +159,8 @@ class _GoogleChunkProcessor:
         self.current_content_type: Literal["text", "tool_call", "thought"] | None = None
         self.accumulated_parts: list[genai_types.Part] = []
         self.reconstructed_content = genai_types.Content(parts=[])
+        # Track previous cumulative usage to compute deltas
+        self.prev_usage = Usage()
 
     def process_chunk(
         self, chunk: genai_types.GenerateContentResponse
@@ -207,6 +236,29 @@ class _GoogleChunkProcessor:
             if finish_reason is not None:
                 yield FinishReasonChunk(finish_reason=finish_reason)
 
+        # Emit usage delta if usage metadata is present
+        if chunk.usage_metadata:
+            usage_metadata = chunk.usage_metadata
+            current_input = usage_metadata.prompt_token_count or 0
+            current_output = usage_metadata.candidates_token_count or 0
+            current_cache_read = usage_metadata.cached_content_token_count or 0
+            current_reasoning = usage_metadata.thoughts_token_count or 0
+
+            yield UsageDeltaChunk(
+                input_tokens=current_input - self.prev_usage.input_tokens,
+                output_tokens=current_output - self.prev_usage.output_tokens,
+                cache_read_tokens=current_cache_read
+                - self.prev_usage.cache_read_tokens,
+                cache_write_tokens=0,
+                reasoning_tokens=current_reasoning - self.prev_usage.reasoning_tokens,
+            )
+
+            # Update previous usage
+            self.prev_usage.input_tokens = current_input
+            self.prev_usage.output_tokens = current_output
+            self.prev_usage.cache_read_tokens = current_cache_read
+            self.prev_usage.reasoning_tokens = current_reasoning
+
     def raw_message_chunk(self) -> RawMessageChunk:
         content = genai_types.Content(role="model", parts=self.accumulated_parts)
         return RawMessageChunk(raw_message=content.model_dump())

mirascope/llm/providers/google/_utils/encode.py

@@ -21,6 +21,7 @@ from ....messages import AssistantMessage, Message, UserMessage
 from ....tools import FORMAT_TOOL_NAME, AnyToolSchema, BaseToolkit
 from ...base import Params, _utils as _base_utils
 from ..model_id import GoogleModelId, model_name
+from ..model_info import MODELS_WITHOUT_STRUCTURED_OUTPUT_AND_TOOLS_SUPPORT
 
 UNKNOWN_TOOL_ID = "google_unknown_tool_id"
 
@@ -187,6 +188,7 @@ def encode_request(
         genai_types.GenerateContentConfigDict()
     )
     encode_thoughts = False
+    google_model_name = model_name(model_id)
 
     with _base_utils.ensure_all_params_accessed(
         params=params, provider_id="google"
@@ -219,17 +221,23 @@ def encode_request(
     tools = tools.tools if isinstance(tools, BaseToolkit) else tools or []
     google_tools: list[genai_types.ToolDict] = []
 
-    format = resolve_format(
-        format,
-        # Google does not support strict outputs when tools are present
-        # (Gemini 2.5 will error, 2.0 and below will ignore tools)
-        default_mode="strict" if not tools else "tool",
+    allows_strict_mode_with_tools = (
+        google_model_name not in MODELS_WITHOUT_STRUCTURED_OUTPUT_AND_TOOLS_SUPPORT
     )
+    # Older google models do not allow strict mode when using tools; if so, we use tool
+    # mode when tools are present by default for compatibility. Otherwise, prefer strict mode.
+    default_mode = "tool" if tools and not allows_strict_mode_with_tools else "strict"
+    format = resolve_format(format, default_mode=default_mode)
     if format is not None:
-        if format.mode in ("strict", "json") and tools:
+        if (
+            format.mode in ("strict", "json")
+            and tools
+            and not allows_strict_mode_with_tools
+        ):
             raise FeatureNotSupportedError(
                 feature=f"formatting_mode:{format.mode} with tools",
                 provider_id="google",
+                model_id=model_id,
             )
 
         if format.mode == "strict":

mirascope/llm/providers/google/_utils/errors.py

@@ -0,0 +1,49 @@
+"""Google error handling utilities."""
+
+from google.genai.errors import (
+    ClientError as GoogleClientError,
+    ServerError as GoogleServerError,
+)
+
+from ....exceptions import (
+    APIError,
+    AuthenticationError,
+    BadRequestError,
+    NotFoundError,
+    PermissionError,
+    RateLimitError,
+    ServerError,
+)
+from ...base import ProviderErrorMap
+
+
+def map_google_error(e: Exception) -> type[APIError]:
+    """Map Google error to appropriate Mirascope error type.
+
+    Google only provides ClientError (4xx) and ServerError (5xx) with status codes,
+    so we map based on status code and message patterns.
+    """
+    if not isinstance(e, GoogleClientError | GoogleServerError):
+        return APIError
+
+    # Authentication errors (401) or 400 with "API key not valid"
+    if e.code == 401 or (e.code == 400 and "API key not valid" in str(e)):
+        return AuthenticationError
+    if e.code == 403:
+        return PermissionError
+    if e.code == 404:
+        return NotFoundError
+    if e.code == 429:
+        return RateLimitError
+    if e.code in (400, 422):
+        return BadRequestError
+    if isinstance(e, GoogleServerError) and e.code >= 500:
+        return ServerError
+    return APIError
+
+
+# Shared error mapping for Google provider
+GOOGLE_ERROR_MAP: ProviderErrorMap = {
+    GoogleClientError: map_google_error,
+    GoogleServerError: map_google_error,
+}

mirascope/llm/providers/google/model_id.py

@@ -1,20 +1,14 @@
 """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
-)
+from typing import TypeAlias, get_args
+
+from .model_info import GoogleKnownModels
+
+GoogleModelId: TypeAlias = GoogleKnownModels | str
 """The Google model ids registered with Mirascope."""
 
+GOOGLE_KNOWN_MODELS: set[str] = set(get_args(GoogleKnownModels))
+
 
 def model_name(model_id: GoogleModelId) -> str:
     """Extract the google model name from a full model ID.

mirascope/llm/providers/google/model_info.py

@@ -0,0 +1,62 @@
+"""Google model information.
+
+This file is auto-generated by scripts/model_features/codegen_google.py
+Do not edit manually - run the codegen script to update."""
+
+from typing import Literal
+
+GoogleKnownModels = Literal[
+    "google/gemini-2.0-flash",
+    "google/gemini-2.0-flash-001",
+    "google/gemini-2.0-flash-exp",
+    "google/gemini-2.0-flash-exp-image-generation",
+    "google/gemini-2.0-flash-lite",
+    "google/gemini-2.0-flash-lite-001",
+    "google/gemini-2.0-flash-lite-preview",
+    "google/gemini-2.0-flash-lite-preview-02-05",
+    "google/gemini-2.5-flash",
+    "google/gemini-2.5-flash-image",
+    "google/gemini-2.5-flash-image-preview",
+    "google/gemini-2.5-flash-lite",
+    "google/gemini-2.5-flash-lite-preview-09-2025",
+    "google/gemini-2.5-flash-preview-09-2025",
+    "google/gemini-2.5-pro",
+    "google/gemini-3-pro-image-preview",
+    "google/gemini-3-pro-preview",
+    "google/gemini-flash-latest",
+    "google/gemini-flash-lite-latest",
+    "google/gemini-pro-latest",
+    "google/gemini-robotics-er-1.5-preview",
+    "google/gemma-3-12b-it",
+    "google/gemma-3-1b-it",
+    "google/gemma-3-27b-it",
+    "google/gemma-3-4b-it",
+    "google/gemma-3n-e2b-it",
+    "google/gemma-3n-e4b-it",
+    "google/nano-banana-pro-preview",
+]
+"""Valid Google model IDs."""
+
+
+MODELS_WITHOUT_STRUCTURED_OUTPUT_AND_TOOLS_SUPPORT: set[str] = {
+    "gemini-2.5-flash",
+    "gemini-2.5-flash-image",
+    "gemini-2.5-flash-image-preview",
+    "gemini-2.5-flash-lite",
+    "gemini-2.5-flash-lite-preview-09-2025",
+    "gemini-2.5-flash-preview-09-2025",
+    "gemini-2.5-pro",
+    "gemini-3-pro-image-preview",
+    "gemini-flash-latest",
+    "gemini-flash-lite-latest",
+    "gemini-pro-latest",
+    "gemini-robotics-er-1.5-preview",
+    "gemma-3-12b-it",
+    "gemma-3-1b-it",
+    "gemma-3-27b-it",
+    "gemma-3-4b-it",
+    "gemma-3n-e2b-it",
+    "gemma-3n-e4b-it",
+    "nano-banana-pro-preview",
+}
+"""Models that do not support structured outputs when tools are present."""