mirascope 2.0.0a2__py3-none-any.whl → 2.0.0a4__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

@@ -29,8 +29,10 @@ from ....responses import (
     FinishReasonChunk,
     RawMessageChunk,
     RawStreamEventChunk,
+    Usage,
+    UsageDeltaChunk,
 )
-from ..model_ids import GoogleModelId
+from ..model_id import GoogleModelId, model_name
 from .encode import UNKNOWN_TOOL_ID
 
 GOOGLE_FINISH_REASON_MAP = {
@@ -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:
@@ -88,7 +114,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,9 +124,10 @@ 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`"""
+    response: genai_types.GenerateContentResponse,
+    model_id: GoogleModelId,
+) -> 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
@@ -115,12 +142,14 @@ 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(),
     )
 
-    return assistant_message, finish_reason
+    usage = _decode_usage(response.usage_metadata)
+    return assistant_message, finish_reason, usage
 
 
 class _GoogleChunkProcessor:
@@ -130,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
@@ -205,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/{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,20 @@ 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
+from ..model_info import MODELS_WITHOUT_STRUCTURED_OUTPUT_AND_TOOLS_SUPPORT
 
 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 +58,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 +127,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 +150,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 +176,22 @@ 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
+    google_model_name = model_name(model_id)
 
     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
@@ -214,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="google",
+                provider_id="google",
+                model_id=model_id,
             )
 
         if format.mode == "strict":
@@ -271,9 +284,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,22 @@
+"""Google registered LLM models."""
+
+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.
+
+    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/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."""