mirascope 2.0.1__py3-none-any.whl → 2.1.0__py3-none-any.whl

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

@@ -27,6 +27,7 @@ from ....responses import (
     ChunkIterator,
     FinishReason,
     FinishReasonChunk,
+    ProviderToolUsage,
     RawMessageChunk,
     RawStreamEventChunk,
     Usage,
@@ -45,8 +46,32 @@ GOOGLE_FINISH_REASON_MAP = {
 }
 
 
+def _extract_tool_usage(
+    candidate: genai_types.Candidate | None,
+) -> list[ProviderToolUsage] | None:
+    """Extract provider tool usage from Google candidate's grounding metadata."""
+    if candidate is None:  # pragma: no cover
+        return None
+
+    grounding_metadata = candidate.grounding_metadata
+    if grounding_metadata is None:
+        return None
+
+    tools: list[ProviderToolUsage] = []
+
+    # Web search queries indicate grounding was used
+    web_search_queries = grounding_metadata.web_search_queries
+    if web_search_queries and len(web_search_queries) > 0:
+        tools.append(
+            ProviderToolUsage(name="web_search", call_count=len(web_search_queries))
+        )
+
+    return tools if tools else None
+
+
 def _decode_usage(
     usage: genai_types.GenerateContentResponseUsageMetadata | None,
+    candidate: genai_types.Candidate | None = None,
 ) -> Usage | None:
     """Convert Google UsageMetadata to Mirascope Usage."""
     if (
@@ -65,6 +90,7 @@ def _decode_usage(
         cache_read_tokens=usage.cached_content_token_count or 0,
         cache_write_tokens=0,
         reasoning_tokens=usage.thoughts_token_count or 0,
+        provider_tool_usage=_extract_tool_usage(candidate),
         raw=usage,
     )
 
@@ -152,7 +178,8 @@ def decode_response(
         raw_message=candidate_content.model_dump(),
     )
 
-    usage = _decode_usage(response.usage_metadata)
+    candidate = response.candidates[0] if response.candidates else None
+    usage = _decode_usage(response.usage_metadata, candidate)
     return assistant_message, finish_reason, usage
 
 
@@ -166,6 +193,8 @@ class _GoogleChunkProcessor:
         # Track previous cumulative usage to compute deltas
         self.prev_usage = Usage()
         self.include_thoughts = include_thoughts
+        # Track web search queries count from grounding metadata
+        self.web_search_query_count = 0
 
     def process_chunk(
         self, chunk: genai_types.GenerateContentResponse
@@ -239,6 +268,15 @@ class _GoogleChunkProcessor:
                 yield ToolCallEndChunk(id=tool_id)
                 self.current_content_type = None
 
+        # Track web search queries from grounding metadata
+        if (
+            candidate.grounding_metadata
+            and candidate.grounding_metadata.web_search_queries
+        ):
+            self.web_search_query_count = len(
+                candidate.grounding_metadata.web_search_queries
+            )
+
         if candidate.finish_reason:
             if self.current_content_type == "text":
                 yield TextEndChunk()
@@ -262,6 +300,17 @@ class _GoogleChunkProcessor:
             current_cache_read = usage_metadata.cached_content_token_count or 0
             current_reasoning = usage_metadata.thoughts_token_count or 0
 
+            # Include provider_tool_usage on the final usage chunk (when finish_reason is present)
+            provider_tool_usage = (
+                [
+                    ProviderToolUsage(
+                        name="web_search", call_count=self.web_search_query_count
+                    )
+                ]
+                if candidate.finish_reason and self.web_search_query_count > 0
+                else None
+            )
+
             yield UsageDeltaChunk(
                 input_tokens=current_input - self.prev_usage.input_tokens,
                 output_tokens=current_output - self.prev_usage.output_tokens,
@@ -269,6 +318,7 @@ class _GoogleChunkProcessor:
                 - self.prev_usage.cache_read_tokens,
                 cache_write_tokens=0,
                 reasoning_tokens=current_reasoning - self.prev_usage.reasoning_tokens,
+                provider_tool_usage=provider_tool_usage,
             )
 
             # Update previous usage

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

@@ -15,12 +15,18 @@ from ....content import ContentPart
 from ....exceptions import FeatureNotSupportedError
 from ....formatting import (
     Format,
+    FormatSpec,
     FormattableT,
-    OutputParser,
     resolve_format,
 )
 from ....messages import AssistantMessage, Message, UserMessage
-from ....tools import FORMAT_TOOL_NAME, AnyToolSchema, BaseToolkit
+from ....tools import (
+    FORMAT_TOOL_NAME,
+    AnyToolSchema,
+    BaseToolkit,
+    ProviderTool,
+    WebSearchTool,
+)
 from ...base import _utils as _base_utils
 from ..model_id import GoogleModelId, model_name
 from ..model_info import MODELS_WITHOUT_STRUCTURED_OUTPUT_AND_TOOLS_SUPPORT
@@ -244,9 +250,13 @@ def _encode_messages(
 
 @lru_cache(maxsize=128)
 def _convert_tool_to_function_declaration(
-    tool: AnyToolSchema,
+    tool: AnyToolSchema | ProviderTool,
 ) -> genai_types.FunctionDeclarationDict:
     """Convert a single Mirascope tool to Google FunctionDeclaration format with caching."""
+    if isinstance(tool, ProviderTool):
+        raise FeatureNotSupportedError(
+            f"Provider tool {tool.name}", provider_id="google"
+        )
     schema_dict = tool.parameters.model_dump(by_alias=True, exclude_none=True)
     schema_dict["type"] = "object"
 
@@ -270,11 +280,8 @@ def encode_request(
     *,
     model_id: GoogleModelId,
     messages: Sequence[Message],
-    tools: Sequence[AnyToolSchema] | BaseToolkit[AnyToolSchema] | None,
-    format: type[FormattableT]
-    | Format[FormattableT]
-    | OutputParser[FormattableT]
-    | None,
+    tools: BaseToolkit[AnyToolSchema],
+    format: FormatSpec[FormattableT] | None,
     params: Params,
 ) -> tuple[Sequence[Message], Format[FormattableT] | None, GoogleKwargs]:
     """Prepares a request for the genai `Client.models.generate_content` method."""
@@ -314,9 +321,7 @@ def encode_request(
             if thinking_config.get("encode_thoughts_as_text"):
                 encode_thoughts_as_text = True
 
-    tools = tools.tools if isinstance(tools, BaseToolkit) else tools or []
-
-    if _base_utils.has_strict_tools(tools):
+    if _base_utils.has_strict_tools(tools.tools):
         raise FeatureNotSupportedError(
             feature="strict tools",
             provider_id="google",
@@ -332,12 +337,14 @@ def encode_request(
     )
     # 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"
+    default_mode = (
+        "tool" if tools.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
+            and tools.tools
             and not allows_strict_mode_with_tools
         ):
             raise FeatureNotSupportedError(
@@ -358,7 +365,7 @@ def encode_request(
             function_calling_config = genai_types.FunctionCallingConfigDict(
                 mode=genai_types.FunctionCallingConfigMode.ANY
             )
-            if not tools:
+            if not tools.tools:
                 function_calling_config["allowed_function_names"] = [FORMAT_TOOL_NAME]
 
             google_config["tool_config"] = genai_types.ToolConfigDict(
@@ -372,13 +379,23 @@ def encode_request(
                 messages, format.formatting_instructions
             )
 
-    if tools:
-        function_declarations = [
-            _convert_tool_to_function_declaration(tool) for tool in tools
-        ]
-        google_tools.append(
-            genai_types.ToolDict(function_declarations=function_declarations)
-        )
+    if tools.tools:
+        # Separate web search tools from function tools
+        function_tools = [t for t in tools.tools if not isinstance(t, WebSearchTool)]
+        has_web_search = any(isinstance(t, WebSearchTool) for t in tools.tools)
+
+        if function_tools:
+            function_declarations = [
+                _convert_tool_to_function_declaration(tool) for tool in function_tools
+            ]
+            google_tools.append(
+                genai_types.ToolDict(function_declarations=function_declarations)
+            )
+
+        if has_web_search:
+            google_tools.append(
+                genai_types.ToolDict(google_search=genai_types.GoogleSearchDict())
+            )
 
     if google_tools:
         google_config["tools"] = cast(genai_types.ToolListUnionDict, google_tools)

mirascope/llm/providers/google/provider.py

@@ -10,7 +10,7 @@ from google.genai import Client
 from google.genai.types import HttpOptions
 
 from ...context import Context, DepsT
-from ...formatting import Format, FormattableT, OutputParser
+from ...formatting import FormatSpec, FormattableT
 from ...messages import Message
 from ...responses import (
     AsyncContextResponse,
@@ -23,13 +23,9 @@ from ...responses import (
     StreamResponse,
 )
 from ...tools import (
-    AsyncContextTool,
     AsyncContextToolkit,
-    AsyncTool,
     AsyncToolkit,
-    ContextTool,
     ContextToolkit,
-    Tool,
     Toolkit,
 )
 from ..base import BaseProvider
@@ -69,11 +65,8 @@ class GoogleProvider(BaseProvider[Client]):
         *,
         model_id: GoogleModelId,
         messages: Sequence[Message],
-        tools: Sequence[Tool] | Toolkit | None = None,
-        format: type[FormattableT]
-        | Format[FormattableT]
-        | OutputParser[FormattableT]
-        | None = None,
+        toolkit: Toolkit,
+        format: FormatSpec[FormattableT] | None = None,
         **params: Unpack[Params],
     ) -> Response | Response[FormattableT]:
         """Generate an `llm.Response` by synchronously calling the Google GenAI API.
@@ -91,7 +84,7 @@ class GoogleProvider(BaseProvider[Client]):
         input_messages, format, kwargs = _utils.encode_request(
             model_id=model_id,
             messages=messages,
-            tools=tools,
+            tools=toolkit,
             format=format,
             params=params,
         )
@@ -108,7 +101,7 @@ class GoogleProvider(BaseProvider[Client]):
             model_id=model_id,
             provider_model_name=model_name(model_id),
             params=params,
-            tools=tools,
+            tools=toolkit,
             input_messages=input_messages,
             assistant_message=assistant_message,
             finish_reason=finish_reason,
@@ -122,13 +115,8 @@ class GoogleProvider(BaseProvider[Client]):
         ctx: Context[DepsT],
         model_id: GoogleModelId,
         messages: Sequence[Message],
-        tools: Sequence[Tool | ContextTool[DepsT]]
-        | ContextToolkit[DepsT]
-        | None = None,
-        format: type[FormattableT]
-        | Format[FormattableT]
-        | OutputParser[FormattableT]
-        | None = None,
+        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.
@@ -147,7 +135,7 @@ class GoogleProvider(BaseProvider[Client]):
         input_messages, format, kwargs = _utils.encode_request(
             model_id=model_id,
             messages=messages,
-            tools=tools,
+            tools=toolkit,
             format=format,
             params=params,
         )
@@ -164,7 +152,7 @@ class GoogleProvider(BaseProvider[Client]):
             model_id=model_id,
             provider_model_name=model_name(model_id),
             params=params,
-            tools=tools,
+            tools=toolkit,
             input_messages=input_messages,
             assistant_message=assistant_message,
             finish_reason=finish_reason,
@@ -177,11 +165,8 @@ class GoogleProvider(BaseProvider[Client]):
         *,
         model_id: GoogleModelId,
         messages: Sequence[Message],
-        tools: Sequence[AsyncTool] | AsyncToolkit | None = None,
-        format: type[FormattableT]
-        | Format[FormattableT]
-        | OutputParser[FormattableT]
-        | None = None,
+        toolkit: AsyncToolkit,
+        format: FormatSpec[FormattableT] | None = None,
         **params: Unpack[Params],
     ) -> AsyncResponse | AsyncResponse[FormattableT]:
         """Generate an `llm.AsyncResponse` by asynchronously calling the Google GenAI API.
@@ -199,7 +184,7 @@ class GoogleProvider(BaseProvider[Client]):
         input_messages, format, kwargs = _utils.encode_request(
             model_id=model_id,
             messages=messages,
-            tools=tools,
+            tools=toolkit,
             format=format,
             params=params,
         )
@@ -216,7 +201,7 @@ class GoogleProvider(BaseProvider[Client]):
             model_id=model_id,
             provider_model_name=model_name(model_id),
             params=params,
-            tools=tools,
+            tools=toolkit,
             input_messages=input_messages,
             assistant_message=assistant_message,
             finish_reason=finish_reason,
@@ -230,13 +215,8 @@ class GoogleProvider(BaseProvider[Client]):
         ctx: Context[DepsT],
         model_id: GoogleModelId,
         messages: Sequence[Message],
-        tools: Sequence[AsyncTool | AsyncContextTool[DepsT]]
-        | AsyncContextToolkit[DepsT]
-        | None = None,
-        format: type[FormattableT]
-        | Format[FormattableT]
-        | OutputParser[FormattableT]
-        | None = None,
+        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.
@@ -255,7 +235,7 @@ class GoogleProvider(BaseProvider[Client]):
         input_messages, format, kwargs = _utils.encode_request(
             model_id=model_id,
             messages=messages,
-            tools=tools,
+            tools=toolkit,
             format=format,
             params=params,
         )
@@ -272,7 +252,7 @@ class GoogleProvider(BaseProvider[Client]):
             model_id=model_id,
             provider_model_name=model_name(model_id),
             params=params,
-            tools=tools,
+            tools=toolkit,
             input_messages=input_messages,
             assistant_message=assistant_message,
             finish_reason=finish_reason,
@@ -285,11 +265,8 @@ class GoogleProvider(BaseProvider[Client]):
         *,
         model_id: GoogleModelId,
         messages: Sequence[Message],
-        tools: Sequence[Tool] | Toolkit | None = None,
-        format: type[FormattableT]
-        | Format[FormattableT]
-        | OutputParser[FormattableT]
-        | None = None,
+        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.
@@ -307,7 +284,7 @@ class GoogleProvider(BaseProvider[Client]):
         input_messages, format, kwargs = _utils.encode_request(
             model_id=model_id,
             messages=messages,
-            tools=tools,
+            tools=toolkit,
             format=format,
             params=params,
         )
@@ -324,7 +301,7 @@ class GoogleProvider(BaseProvider[Client]):
             model_id=model_id,
             provider_model_name=model_name(model_id),
             params=params,
-            tools=tools,
+            tools=toolkit,
             input_messages=input_messages,
             chunk_iterator=chunk_iterator,
             format=format,
@@ -336,13 +313,8 @@ class GoogleProvider(BaseProvider[Client]):
         ctx: Context[DepsT],
         model_id: GoogleModelId,
         messages: Sequence[Message],
-        tools: Sequence[Tool | ContextTool[DepsT]]
-        | ContextToolkit[DepsT]
-        | None = None,
-        format: type[FormattableT]
-        | Format[FormattableT]
-        | OutputParser[FormattableT]
-        | None = None,
+        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.
@@ -361,7 +333,7 @@ class GoogleProvider(BaseProvider[Client]):
         input_messages, format, kwargs = _utils.encode_request(
             model_id=model_id,
             messages=messages,
-            tools=tools,
+            tools=toolkit,
             format=format,
             params=params,
         )
@@ -378,7 +350,7 @@ class GoogleProvider(BaseProvider[Client]):
             model_id=model_id,
             provider_model_name=model_name(model_id),
             params=params,
-            tools=tools,
+            tools=toolkit,
             input_messages=input_messages,
             chunk_iterator=chunk_iterator,
             format=format,
@@ -389,11 +361,8 @@ class GoogleProvider(BaseProvider[Client]):
         *,
         model_id: GoogleModelId,
         messages: Sequence[Message],
-        tools: Sequence[AsyncTool] | AsyncToolkit | None = None,
-        format: type[FormattableT]
-        | Format[FormattableT]
-        | OutputParser[FormattableT]
-        | None = None,
+        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.
@@ -411,7 +380,7 @@ class GoogleProvider(BaseProvider[Client]):
         input_messages, format, kwargs = _utils.encode_request(
             model_id=model_id,
             messages=messages,
-            tools=tools,
+            tools=toolkit,
             format=format,
             params=params,
         )
@@ -428,7 +397,7 @@ class GoogleProvider(BaseProvider[Client]):
             model_id=model_id,
             provider_model_name=model_name(model_id),
             params=params,
-            tools=tools,
+            tools=toolkit,
             input_messages=input_messages,
             chunk_iterator=chunk_iterator,
             format=format,
@@ -440,13 +409,8 @@ class GoogleProvider(BaseProvider[Client]):
         ctx: Context[DepsT],
         model_id: GoogleModelId,
         messages: Sequence[Message],
-        tools: Sequence[AsyncTool | AsyncContextTool[DepsT]]
-        | AsyncContextToolkit[DepsT]
-        | None = None,
-        format: type[FormattableT]
-        | Format[FormattableT]
-        | OutputParser[FormattableT]
-        | None = None,
+        toolkit: AsyncContextToolkit[DepsT],
+        format: FormatSpec[FormattableT] | None = None,
         **params: Unpack[Params],
     ) -> (
         AsyncContextStreamResponse[DepsT]
@@ -468,7 +432,7 @@ class GoogleProvider(BaseProvider[Client]):
         input_messages, format, kwargs = _utils.encode_request(
             model_id=model_id,
             messages=messages,
-            tools=tools,
+            tools=toolkit,
             format=format,
             params=params,
         )
@@ -485,7 +449,7 @@ class GoogleProvider(BaseProvider[Client]):
             model_id=model_id,
             provider_model_name=model_name(model_id),
             params=params,
-            tools=tools,
+            tools=toolkit,
             input_messages=input_messages,
             chunk_iterator=chunk_iterator,
             format=format,