mirascope 1.16.9__py3-none-any.whl → 1.18.0__py3-none-any.whl

mirascope/core/gemini/_utils/_message_param_converter.py

@@ -8,7 +8,14 @@ from google.generativeai.types import (
 )
 
 from mirascope.core import BaseMessageParam
-from mirascope.core.base import DocumentPart, ImagePart, TextPart
+from mirascope.core.base import (
+    AudioPart,
+    AudioURLPart,
+    DocumentPart,
+    ImagePart,
+    ImageURLPart,
+    TextPart,
+)
 from mirascope.core.base._utils._base_message_param_converter import (
     BaseMessageParamConverter,
 )
@@ -29,6 +36,28 @@ def _to_image_part(mime_type: str, data: bytes) -> ImagePart:
     return ImagePart(type="image", media_type=mime_type, image=data, detail=None)
 
 
+def _is_audio_mime(mime_type: str) -> bool:
+    return mime_type in [
+        "audio/wav",
+        "audio/mp3",
+        "audio/wav",
+        "audio/mp3",
+        "audio/aiff",
+        "audio/aac",
+        "audio/ogg",
+        "audio/flac",
+    ]
+
+
+def _to_audio_part(mime_type: str, data: bytes) -> AudioPart:
+    if not _is_audio_mime(mime_type):
+        raise ValueError(
+            f"Unsupported audio media type: {mime_type}. "
+            "Expected one of: audio/wav, audio/mp3, audio/aiff, audio/aac, audio/ogg, audio/flac."
+        )
+    return AudioPart(type="audio", media_type=mime_type, audio=data)
+
+
 def _to_document_part(mime_type: str, data: bytes) -> DocumentPart:
     if mime_type != "application/pdf":
         raise ValueError(
@@ -77,6 +106,8 @@ class GeminiMessageParamConverter(BaseMessageParamConverter):
                         data = blob.data
                         if _is_image_mime(mime):
                             content_list.append(_to_image_part(mime, data))
+                        elif _is_audio_mime(mime):
+                            content_list.append(_to_audio_part(mime, data))
                         elif mime == "application/pdf":
                             content_list.append(_to_document_part(mime, data))
                         else:
@@ -85,10 +116,25 @@ class GeminiMessageParamConverter(BaseMessageParamConverter):
                             )
 
                     elif part.file_data:
-                        # part.file_data.file_uri has Google storage URI like "gs://bucket_name/file_name"
-                        raise ValueError(
-                            f"FileData.file_uri is not support: {part.file_data}. Cannot convert to BaseMessageParam."
-                        )
+                        if _is_image_mime(part.file_data.mime_type):
+                            content_list.append(
+                                ImageURLPart(
+                                    type="image_url",
+                                    url=part.file_data.file_uri,
+                                    detail=None,
+                                )
+                            )
+                        elif _is_audio_mime(part.file_data.mime_type):
+                            content_list.append(
+                                AudioURLPart(
+                                    type="audio_url",
+                                    url=part.file_data.file_uri,
+                                )
+                            )
+                        else:
+                            raise ValueError(
+                                f"Unsupported file_data mime type: {part.file_data.mime_type}. Cannot convert to BaseMessageParam."
+                            )
                     elif part.function_call:
                         converted.append(
                             BaseMessageParam(

mirascope/core/gemini/_utils/_setup_call.py

@@ -11,7 +11,7 @@ from google.generativeai.types import (
     GenerateContentResponse,
     GenerationConfigDict,
 )
-from google.generativeai.types.content_types import ToolConfigDict
+from google.generativeai.types.content_types import ToolConfigDict, to_content
 from pydantic import BaseModel
 
 from ...base import BaseMessageParam, BaseTool, _utils
@@ -108,6 +108,7 @@ def setup_call(
     call_kwargs = cast(GeminiCallKwargs, base_call_kwargs)
     messages = cast(list[BaseMessageParam | ContentDict], messages)
     messages = convert_message_params(messages)
+
     if json_mode:
         generation_config = call_kwargs.get("generation_config", {})
         if is_dataclass(generation_config):
@@ -125,11 +126,20 @@ def setup_call(
             "allowed_function_names": [tool_types[0]._name()],
         }
         call_kwargs["tool_config"] = tool_config
-    call_kwargs |= {"contents": messages}
 
     if client is None:
         client = GenerativeModel(model_name=model)
 
+    if messages and messages[0]["role"] == "system":
+        system_instruction = client._system_instruction
+        system_instruction = (
+            list(system_instruction.parts) if system_instruction else []
+        )
+        system_instruction.extend(messages.pop(0)["parts"])  # pyright: ignore [reportArgumentType]
+        client._system_instruction = to_content(system_instruction)  # pyright: ignore [reportArgumentType]
+
+    call_kwargs |= {"contents": messages}
+
     create = (
         get_async_create_fn(client.generate_content_async)
         if fn_is_async(fn)

mirascope/core/google/__init__.py

@@ -0,0 +1,29 @@
+"""The Mirascope Google Module."""
+
+from typing import TypeAlias
+
+from google.genai.types import ContentDict, FunctionResponse
+
+from ..base import BaseMessageParam
+from ._call import google_call
+from ._call import google_call as call
+from .call_params import GoogleCallParams
+from .call_response import GoogleCallResponse
+from .call_response_chunk import GoogleCallResponseChunk
+from .dynamic_config import GoogleDynamicConfig
+from .stream import GoogleStream
+from .tool import GoogleTool
+
+GoogleMessageParam: TypeAlias = ContentDict | FunctionResponse | BaseMessageParam
+
+__all__ = [
+    "call",
+    "GoogleDynamicConfig",
+    "GoogleCallParams",
+    "GoogleCallResponse",
+    "GoogleCallResponseChunk",
+    "GoogleMessageParam",
+    "GoogleStream",
+    "GoogleTool",
+    "google_call",
+]

mirascope/core/google/_call.py

@@ -0,0 +1,67 @@
+"""The `google_call` decorator for functions as LLM calls."""
+
+from ..base import call_factory
+from ._utils import (
+    get_json_output,
+    handle_stream,
+    handle_stream_async,
+    setup_call,
+)
+from .call_params import GoogleCallParams
+from .call_response import GoogleCallResponse
+from .call_response_chunk import GoogleCallResponseChunk
+from .stream import GoogleStream
+from .tool import GoogleTool
+
+google_call = call_factory(
+    TCallResponse=GoogleCallResponse,
+    TCallResponseChunk=GoogleCallResponseChunk,
+    TStream=GoogleStream,
+    TToolType=GoogleTool,
+    default_call_params=GoogleCallParams(),
+    setup_call=setup_call,
+    get_json_output=get_json_output,
+    handle_stream=handle_stream,  # pyright: ignore [reportArgumentType]
+    handle_stream_async=handle_stream_async,  # pyright: ignore [reportArgumentType]
+)
+"""A decorator for calling the Google API with a typed function.
+
+usage docs: learn/calls.md
+
+This decorator is used to wrap a typed function that calls the Google API. It parses
+the prompt template of the wrapped function as the messages array and templates the input
+arguments for the function into each message's template.
+
+Example:
+
+```python
+from mirascope.core import prompt_template
+from mirascope.core.google import google_call
+
+
+@google_call("google-1.5-flash")
+def recommend_book(genre: str) -> str:
+    return f"Recommend a {genre} book"
+
+response = recommend_book("fantasy")
+print(response.content)
+```
+
+Args:
+    model (str): The Google model to use in the API call.
+    stream (bool): Whether to stream the response from the API call.
+    tools (list[BaseTool | Callable]): The tools to use in the Google API call.
+    response_model (BaseModel | BaseType): The response model into which the response
+        should be structured.
+    output_parser (Callable[[GoogleCallResponse | ResponseModelT], Any]): A function
+        for parsing the call response whose value will be returned in place of the
+        original call response.
+    json_modem (bool): Whether to use JSON Mode.
+    client (object): An optional custom client to use in place of the default client.
+    call_params (GoogleCallParams): The `GoogleCallParams` call parameters to use in the
+        API call.
+
+Returns:
+    decorator (Callable): The decorator for turning a typed function into a Google API
+        call.
+"""

mirascope/core/google/_call_kwargs.py

@@ -0,0 +1,13 @@
+"""This module contains the type definition for the Google call keyword arguments."""
+
+from collections.abc import Sequence
+
+from google.genai.types import ContentOrDict, Tool
+
+from ..base import BaseCallKwargs
+from .call_params import GoogleCallParams
+
+
+class GoogleCallKwargs(GoogleCallParams, BaseCallKwargs[Tool]):
+    model: str
+    contents: Sequence[ContentOrDict]

mirascope/core/google/_utils/__init__.py

@@ -0,0 +1,16 @@
+"""Google utilities for decorator factories."""
+
+from ._calculate_cost import calculate_cost
+from ._convert_message_params import convert_message_params
+from ._get_json_output import get_json_output
+from ._handle_stream import handle_stream, handle_stream_async
+from ._setup_call import setup_call
+
+__all__ = [
+    "calculate_cost",
+    "convert_message_params",
+    "get_json_output",
+    "handle_stream",
+    "handle_stream_async",
+    "setup_call",
+]

mirascope/core/google/_utils/_calculate_cost.py

@@ -0,0 +1,88 @@
+"""Calculate the cost of a Gemini API call."""
+
+
+def calculate_cost(
+    input_tokens: int | float | None, output_tokens: int | float | None, model: str
+) -> float | None:
+    """Calculate the cost of a Google API call.
+
+    https://ai.google.dev/pricing
+
+    Pricing (per 1M tokens):
+
+    Model                   Input (<128K)    Output (<128K)    Input (>128K)    Output (>128K)
+    gemini-2.0-flash       $0.10            $0.40            $0.10            $0.40
+    gemini-2.0-flash-lite  $0.075           $0.30            $0.075           $0.30
+    gemini-1.5-flash       $0.075           $0.30            $0.15            $0.60
+    gemini-1.5-flash-8b    $0.0375          $0.15            $0.075           $0.30
+    gemini-1.5-pro         $1.25            $5.00            $2.50            $10.00
+    gemini-1.0-pro         $0.50            $1.50            $0.50            $1.50
+
+    Args:
+        input_tokens: Number of input tokens
+        output_tokens: Number of output tokens
+        model: Model name to use for pricing calculation
+
+    Returns:
+        Total cost in USD or None if invalid input
+    """
+    pricing = {
+        "gemini-2.0-flash": {
+            "prompt_short": 0.000_000_10,
+            "completion_short": 0.000_000_40,
+            "prompt_long": 0.000_000_10,
+            "completion_long": 0.000_000_40,
+        },
+        "gemini-2.0-flash-lite": {
+            "prompt_short": 0.000_000_075,
+            "completion_short": 0.000_000_30,
+            "prompt_long": 0.000_000_075,
+            "completion_long": 0.000_000_30,
+        },
+        "gemini-1.5-flash": {
+            "prompt_short": 0.000_000_075,
+            "completion_short": 0.000_000_30,
+            "prompt_long": 0.000_000_15,
+            "completion_long": 0.000_000_60,
+        },
+        "gemini-1.5-flash-8b": {
+            "prompt_short": 0.000_000_037_5,
+            "completion_short": 0.000_000_15,
+            "prompt_long": 0.000_000_075,
+            "completion_long": 0.000_000_30,
+        },
+        "gemini-1.5-pro": {
+            "prompt_short": 0.000_001_25,
+            "completion_short": 0.000_005,
+            "prompt_long": 0.000_002_5,
+            "completion_long": 0.000_01,
+        },
+        "gemini-1.0-pro": {
+            "prompt_short": 0.000_000_5,
+            "completion_short": 0.000_001_5,
+            "prompt_long": 0.000_000_5,
+            "completion_long": 0.000_001_5,
+        },
+    }
+
+    if input_tokens is None or output_tokens is None:
+        return None
+
+    try:
+        model_pricing = pricing[model]
+    except KeyError:
+        return None
+
+    # Determine if we're using long context pricing
+    use_long_context = input_tokens > 128_000
+
+    prompt_price = model_pricing["prompt_long" if use_long_context else "prompt_short"]
+    completion_price = model_pricing[
+        "completion_long" if use_long_context else "completion_short"
+    ]
+
+    prompt_cost = input_tokens * prompt_price
+    completion_cost = output_tokens * completion_price
+    total_cost = prompt_cost + completion_cost
+
+    return total_cost

mirascope/core/google/_utils/_convert_common_call_params.py

@@ -0,0 +1,39 @@
+from typing import cast
+
+from google.genai.types import (
+    GenerationConfigDict,
+)
+
+from ...base.call_params import CommonCallParams
+from ..call_params import GoogleCallParams
+
+GOOGLE_PARAM_MAPPING = {
+    "temperature": "temperature",
+    "max_tokens": "max_output_tokens",
+    "top_p": "top_p",
+    "stop": "stop_sequences",
+}
+
+
+def convert_common_call_params(common_params: CommonCallParams) -> GoogleCallParams:
+    """Convert CommonCallParams to Google parameters."""
+    generation_config = {}
+
+    for key, value in common_params.items():
+        if key not in GOOGLE_PARAM_MAPPING or value is None:
+            continue
+
+        if key == "stop":
+            generation_config["stop_sequences"] = (
+                [value] if isinstance(value, str) else value
+            )
+        else:
+            generation_config[GOOGLE_PARAM_MAPPING[key]] = value
+
+    if not generation_config:
+        return cast(GoogleCallParams, {})
+
+    return cast(
+        GoogleCallParams,
+        {"generation_config": cast(GenerationConfigDict, generation_config)},
+    )

mirascope/core/google/_utils/_convert_finish_reason_to_common_finish_reasons.py

@@ -0,0 +1,27 @@
+from mirascope.core.base._utils._convert_provider_finish_reason_to_finish_reason import (
+    FinishReasonMappingValue,
+    _convert_finish_reasons_to_common_finish_reasons_from_mapping,
+)
+from mirascope.core.base.types import FinishReason
+
+_FinishReasonMapping: dict[str, FinishReasonMappingValue] = {
+    "FINISH_REASON_UNSPECIFIED": "stop",
+    "STOP": "stop",
+    "MAX_TOKENS": "length",
+    "SAFETY": "content_filter",
+    "RECITATION": "stop",
+    "OTHER": "stop",
+    "BLOCKLIST": "stop",
+    "PROHIBITED_CONTENT": "stop",
+    "SPII": "stop",
+    "MALFORMED_FUNCTION_CALL": "stop",
+}
+
+
+def _convert_finish_reasons_to_common_finish_reasons(
+    finish_reasons: list[str],
+) -> list[FinishReason] | None:
+    """Provider-agnostic finish reasons."""
+    return _convert_finish_reasons_to_common_finish_reasons_from_mapping(
+        finish_reasons, _FinishReasonMapping
+    )

mirascope/core/google/_utils/_convert_message_params.py

@@ -0,0 +1,177 @@
+"""Utility for converting `BaseMessageParam` to `ContentsType`"""
+
+import base64
+import io
+
+import PIL.Image
+from google.genai import Client
+from google.genai.types import BlobDict, ContentDict, FileDataDict, PartDict
+
+from ...base import BaseMessageParam
+from ...base._utils import get_audio_type
+from ...base._utils._parse_content_template import _load_media
+
+
+def convert_message_params(
+    message_params: list[BaseMessageParam | ContentDict], client: Client
+) -> list[ContentDict]:
+    converted_message_params = []
+    for message_param in message_params:
+        if not isinstance(message_param, BaseMessageParam):
+            converted_message_params.append(message_param)
+        elif (role := message_param.role) == "system":
+            if not isinstance(message_param.content, str):
+                raise ValueError(
+                    "System message content must be a single text string."
+                )  # pragma: no cover
+            converted_message_params += [
+                {
+                    "role": "system",
+                    "parts": [PartDict(text=message_param.content)],
+                }
+            ]
+        elif isinstance((content := message_param.content), str):
+            converted_message_params.append(
+                {
+                    "role": role if role == "user" else "model",
+                    "parts": [PartDict(text=content)],
+                }
+            )
+        else:
+            converted_content = []
+            for part in content:
+                if part.type == "text":
+                    converted_content.append(PartDict(text=part.text))
+                elif part.type == "image":
+                    if part.media_type not in [
+                        "image/jpeg",
+                        "image/png",
+                        "image/webp",
+                        "image/heic",
+                        "image/heif",
+                    ]:
+                        raise ValueError(
+                            f"Unsupported image media type: {part.media_type}. "
+                            "Google currently only supports JPEG, PNG, WebP, HEIC, "
+                            "and HEIF images."
+                        )
+                    converted_content.append(
+                        PartDict(
+                            inline_data=BlobDict(
+                                data=part.image, mime_type=part.media_type
+                            )
+                        )
+                    )
+                elif part.type == "image_url":
+                    if part.url.startswith(("https://", "http://")):
+                        downloaded_image = io.BytesIO(_load_media(part.url))
+                        image = PIL.Image.open(downloaded_image)
+                        media_type = (
+                            PIL.Image.MIME[image.format]
+                            if image.format
+                            else "image/unknown"
+                        )
+                        if media_type not in [
+                            "image/jpeg",
+                            "image/png",
+                            "image/webp",
+                            "image/heic",
+                            "image/heif",
+                        ]:
+                            raise ValueError(
+                                f"Unsupported image media type: {media_type}. "
+                                "Google currently only supports JPEG, PNG, WebP, HEIC, "
+                                "and HEIF images."
+                            )
+                        if client.vertexai:
+                            uri = part.url
+                        else:
+                            downloaded_image.seek(0)
+                            file_ref = client.files.upload(
+                                file=downloaded_image, config={"mime_type": media_type}
+                            )
+                            uri = file_ref.uri
+                            media_type = file_ref.mime_type
+
+                        converted_content.append(
+                            PartDict(
+                                file_data=FileDataDict(
+                                    file_uri=uri, mime_type=media_type
+                                )
+                            )
+                        )
+                    else:
+                        media_type = "image/unknown"
+                        uri = part.url
+                        converted_content.append(
+                            PartDict(
+                                file_data=FileDataDict(
+                                    file_uri=uri, mime_type=media_type
+                                )
+                            )
+                        )
+                elif part.type == "audio":
+                    if part.media_type not in [
+                        "audio/wav",
+                        "audio/mp3",
+                        "audio/aiff",
+                        "audio/aac",
+                        "audio/ogg",
+                        "audio/flac",
+                    ]:
+                        raise ValueError(
+                            f"Unsupported audio media type: {part.media_type}. "
+                            "Google currently only supports WAV, MP3, AIFF, AAC, OGG, "
+                            "and FLAC audio file types."
+                        )
+                    converted_content.append(
+                        PartDict(
+                            inline_data=BlobDict(
+                                data=part.audio
+                                if isinstance(part.audio, bytes)
+                                else base64.b64decode(part.audio),
+                                mime_type=part.media_type,
+                            )
+                        )
+                    )
+                elif part.type == "audio_url":
+                    if part.url.startswith(("https://", "http://")):
+                        audio = _load_media(part.url)
+                        audio_type = get_audio_type(audio)
+                        if audio_type not in [
+                            "audio/wav",
+                            "audio/mp3",
+                            "audio/aiff",
+                            "audio/aac",
+                            "audio/ogg",
+                            "audio/flac",
+                        ]:
+                            raise ValueError(
+                                f"Unsupported audio media type: {audio_type}. "
+                                "Google currently only supports WAV, MP3, AIFF, AAC, OGG, "
+                                "and FLAC audio file types."
+                            )
+                        converted_content.append(
+                            {"mime_type": audio_type, "data": audio}
+                        )
+                    else:
+                        audio_type = "audio/unknown"
+                    converted_content.append(
+                        PartDict(
+                            file_data=FileDataDict(
+                                file_uri=part.url, mime_type=audio_type
+                            )
+                        )
+                    )
+                else:
+                    raise ValueError(
+                        "Google currently only supports text, image, and audio parts. "
+                        f"Part provided: {part.type}"
+                    )
+            converted_message_params.append(
+                {
+                    "role": role if role == "user" else "model",
+                    "parts": converted_content,
+                }
+            )
+    return converted_message_params

mirascope/core/google/_utils/_get_json_output.py

@@ -0,0 +1,37 @@
+"""Get JSON output from a Google response."""
+
+import json
+
+from proto.marshal.collections import RepeatedComposite
+
+from ..call_response import GoogleCallResponse
+from ..call_response_chunk import GoogleCallResponseChunk
+
+
+def get_json_output(
+    response: GoogleCallResponse | GoogleCallResponseChunk, json_mode: bool
+) -> str:
+    """Extracts the JSON output from a Google response."""
+    if isinstance(response, GoogleCallResponse):
+        if json_mode and (content := response.content):
+            json_start = content.index("{")
+            json_end = content.rfind("}")
+            return content[json_start : json_end + 1]
+        elif tool_calls := [
+            function_call
+            for function_call in (response.response.function_calls or [])  # pyright: ignore [reportOptionalSubscript, reportOptionalIterable, reportOptionalMemberAccess]
+            if function_call.args  # pyright: ignore [reportOptionalSubscript, reportOptionalIterable, reportOptionalMemberAccess]
+        ]:
+            return json.dumps(
+                {
+                    k: v if not isinstance(v, RepeatedComposite) else list(v)
+                    for k, v in (tool_calls[0].args or {}).items()  # pyright: ignore [reportOptionalMemberAccess]
+                }
+                if isinstance(tool_calls, list) and tool_calls[0]
+                else {}
+            )
+        else:
+            raise ValueError("No tool call or JSON object found in response.")
+    elif not json_mode:
+        raise ValueError("Google only supports structured streaming in json mode.")
+    return response.content

mirascope/core/google/_utils/_handle_stream.py

@@ -0,0 +1,35 @@
+"""Handles streaming content and tools from the Google API."""
+
+from collections.abc import AsyncGenerator, Generator
+
+from google.genai.types import GenerateContentResponse
+
+from ..call_response_chunk import GoogleCallResponseChunk
+from ..tool import GoogleTool
+
+
+def handle_stream(
+    stream: Generator[GenerateContentResponse, None, None],
+    tool_types: list[type[GoogleTool]] | None = None,
+    partial_tools: bool = False,
+) -> Generator[tuple[GoogleCallResponseChunk, None], None, None]:
+    """Iterator over the stream and constructs tools as they are streamed.
+
+    Note: google does not currently support streaming tools.
+    """
+    for chunk in stream:
+        yield GoogleCallResponseChunk(chunk=chunk), None
+
+
+async def handle_stream_async(
+    stream: AsyncGenerator[GenerateContentResponse, None],
+    tool_types: list[type[GoogleTool]] | None = None,
+    partial_tools: bool = False,
+) -> AsyncGenerator[tuple[GoogleCallResponseChunk, None], None]:
+    """
+    Async iterator over the stream and constructs tools as they are streamed.
+
+    Note: google does not currently support streaming tools.
+    """
+    async for chunk in stream:
+        yield GoogleCallResponseChunk(chunk=chunk), None