Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions src/celeste/providers/google/generate_content/parameters.py
Original file line number Diff line number Diff line change
Expand Up @@ -184,9 +184,11 @@ def map(
dispatch = {m.tool_type: m for m in TOOL_MAPPERS}
tools = request.setdefault("tools", [])
fn_declarations: list[dict[str, Any]] = []
has_server_side_tool = False

for item in validated_value:
if isinstance(item, Tool):
has_server_side_tool = True
mapper = dispatch.get(type(item))
if mapper is None:
msg = f"Tool '{type(item).__name__}' is not supported by Google"
Expand All @@ -195,13 +197,20 @@ def map(
elif isinstance(item, dict) and "name" in item:
fn_declarations.append(self._map_user_tool(item))
elif isinstance(item, dict):
has_server_side_tool = True
tools.append(item)
else:
raise InvalidToolError(item)

if fn_declarations:
tools.append({"functionDeclarations": fn_declarations})

# Gemini requires this when mixing built-in tools with functionDeclarations.
if fn_declarations and has_server_side_tool:
request.setdefault("toolConfig", {})["includeServerSideToolInvocations"] = (
True
)

return request

@staticmethod
Expand Down
49 changes: 48 additions & 1 deletion tests/unit_tests/test_google_tools_mapper.py
Original file line number Diff line number Diff line change
@@ -1,14 +1,38 @@
"""Regression coverage for Gemini ToolsMapper — emits `parametersJsonSchema`."""
"""Regression coverage for Gemini ToolsMapper."""

from __future__ import annotations

from typing import Any

import pytest

from celeste.core import Provider
from celeste.modalities.text.parameters import TextParameter
from celeste.models import Model
from celeste.providers.google.generate_content.parameters import ToolsMapper
from celeste.tools import CodeExecution, ToolDefinition, WebSearch


class _GoogleToolsMapper(ToolsMapper):
name = TextParameter.TOOLS


_MAPPER = _GoogleToolsMapper()
_MODEL = Model(id="test-model", provider=Provider.GOOGLE, display_name="Test Model")
_FUNCTION_TOOL: dict[str, Any] = {
"name": "web_fetch",
"parameters": {"type": "object"},
}


def _map(tool: dict) -> dict:
return ToolsMapper._map_user_tool(tool)


def _map_tools(tools: list[ToolDefinition]) -> dict[str, Any]:
return _MAPPER.map({}, tools, _MODEL)


def test_maps_to_parameters_json_schema_not_legacy_parameters() -> None:
fn = _map(
{
Expand Down Expand Up @@ -63,3 +87,26 @@ def test_tool_with_no_parameters_emits_no_schema_field() -> None:
assert "parameters" not in fn
assert "parametersJsonSchema" not in fn
assert fn["description"] == "no-args"


@pytest.mark.parametrize(
("tools", "expect_flag"),
[
([WebSearch(), _FUNCTION_TOOL], True),
([CodeExecution(), _FUNCTION_TOOL], True),
([{"google_search": {}}, _FUNCTION_TOOL], True),
([_FUNCTION_TOOL], False),
([WebSearch()], False),
([CodeExecution()], False),
],
)
def test_include_server_side_tool_invocations(
tools: list[ToolDefinition], expect_flag: bool
) -> None:
result = _map_tools(tools)
flag = result.get("toolConfig", {}).get("includeServerSideToolInvocations")
if expect_flag:
assert flag is True
assert any("functionDeclarations" in tool for tool in result["tools"])
else:
assert flag is None
Loading