mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-22 13:34:17 +02:00
Add Google GenAI Sampling Handler (#2977)
* Add Google GenAI sampling handler Co-authored-by: Bill Easton <strawgate@users.noreply.github.com> * chore: Update SDK documentation * Filter non-Gemini model hints in _get_model Match the Anthropic/OpenAI handler pattern of only selecting provider-compatible models from hints. --------- Co-authored-by: claude[bot] <41898282+claude[bot]@users.noreply.github.com> Co-authored-by: Bill Easton <strawgate@users.noreply.github.com> Co-authored-by: marvin-context-protocol[bot] <225465937+marvin-context-protocol[bot]@users.noreply.github.com> Co-authored-by: Jeremiah Lowin <153965+jlowin@users.noreply.github.com>
This commit is contained in:
Bill Easton
GitHub
parent
af557388c7
commit
49534ce91d
6 changed files with 868 additions and 4 deletions
|
|
@ -23,6 +23,38 @@ result = await ctx.sample(
|
|||
)
|
||||
```
|
||||
|
||||
### Google GenAI Sampling Handler
|
||||
|
||||
FastMCP now includes a sampling handler for Google's Gemini models ([#2977](https://github.com/jlowin/fastmcp/pull/2977)). This enables MCP clients to use Google's GenAI models with the sampling protocol, including full tool calling support.
|
||||
|
||||
```python
|
||||
from fastmcp import Client
|
||||
from fastmcp.client.sampling.handlers import GoogleGenaiSamplingHandler
|
||||
from google.genai import Client as GoogleGenaiClient
|
||||
|
||||
# Initialize the handler
|
||||
handler = GoogleGenaiSamplingHandler(
|
||||
default_model="gemini-2.0-flash-exp",
|
||||
client=GoogleGenaiClient(), # Optional - creates one if not provided
|
||||
)
|
||||
|
||||
# Use with MCP sampling (handler is configured at Client construction)
|
||||
async with Client("http://server/mcp", sampling_handler=handler) as client:
|
||||
result = await client.sample(
|
||||
messages=[...],
|
||||
tools=[...],
|
||||
)
|
||||
```
|
||||
|
||||
Key features:
|
||||
- Converts MCP tool schemas to Google's function calling format
|
||||
- Supports all Google GenAI models that implement function calling
|
||||
- Handles nullable types, nested objects, and arrays in tool schemas
|
||||
- Properly maps tool choices (`auto`, `required`, `none`) to Google's configuration
|
||||
- Preserves model preferences from MCP sampling parameters
|
||||
|
||||
The handler joins the existing Anthropic and OpenAI handlers, providing a consistent interface for model-agnostic sampling across providers.
|
||||
|
||||
### Concurrent Tool Execution in Sampling
|
||||
|
||||
When an LLM returns multiple tool calls in a single sampling response, they can now be executed concurrently ([#3022](https://github.com/PrefectHQ/fastmcp/pull/3022)). Default behavior remains sequential; opt in with `tool_concurrency`. Tools can declare `sequential=True` to force sequential execution even when concurrency is enabled.
|
||||
|
|
|
|||
|
|
@ -0,0 +1,17 @@
|
|||
---
|
||||
title: google_genai
|
||||
sidebarTitle: google_genai
|
||||
---
|
||||
|
||||
# `fastmcp.client.sampling.handlers.google_genai`
|
||||
|
||||
|
||||
Google GenAI sampling handler with tool support for FastMCP 3.0.
|
||||
|
||||
## Classes
|
||||
|
||||
### `GoogleGenaiSamplingHandler` <sup><a href="https://github.com/jlowin/fastmcp/blob/main/src/fastmcp/client/sampling/handlers/google_genai.py#L54" target="_blank"><Icon icon="github" style="width: 14px; height: 14px;" /></a></sup>
|
||||
|
||||
|
||||
Sampling handler that uses the Google GenAI API with tool support.
|
||||
|
||||
|
|
@ -56,13 +56,14 @@ anthropic = ["anthropic>=0.40.0"]
|
|||
apps = ["prefab-ui>=0.6.0"]
|
||||
azure = ["azure-identity>=1.16.0"]
|
||||
code-mode = ["pydantic-monty>=0.0.7"]
|
||||
gemini = ["google-genai>=1.18.0"]
|
||||
openai = ["openai>=1.102.0"]
|
||||
tasks = ["pydocket>=0.18.0"]
|
||||
|
||||
[dependency-groups]
|
||||
dev = [
|
||||
"dirty-equals>=0.9.0",
|
||||
"fastmcp[anthropic,apps,azure,code-mode,openai,tasks]",
|
||||
"fastmcp[anthropic,apps,azure,code-mode,gemini,openai,tasks]",
|
||||
# add optional dependencies for fastmcp dev
|
||||
"fastapi>=0.115.12",
|
||||
"opentelemetry-sdk>=1.20.0",
|
||||
|
|
|
|||
368
src/fastmcp/client/sampling/handlers/google_genai.py
Normal file
368
src/fastmcp/client/sampling/handlers/google_genai.py
Normal file
|
|
@ -0,0 +1,368 @@
|
|||
"""Google GenAI sampling handler with tool support for FastMCP 3.0."""
|
||||
|
||||
from collections.abc import Sequence
|
||||
from uuid import uuid4
|
||||
|
||||
try:
|
||||
from google.genai import Client as GoogleGenaiClient
|
||||
from google.genai.types import (
|
||||
Candidate,
|
||||
Content,
|
||||
FunctionCall,
|
||||
FunctionCallingConfig,
|
||||
FunctionCallingConfigMode,
|
||||
FunctionDeclaration,
|
||||
FunctionResponse,
|
||||
GenerateContentConfig,
|
||||
GenerateContentResponse,
|
||||
ModelContent,
|
||||
Part,
|
||||
ThinkingConfig,
|
||||
ToolConfig,
|
||||
UserContent,
|
||||
)
|
||||
from google.genai.types import Tool as GoogleTool
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"The `google-genai` package is not installed. "
|
||||
"Install it with `pip install fastmcp[gemini]` or add `google-genai` "
|
||||
"to your dependencies."
|
||||
) from e
|
||||
|
||||
from mcp import ClientSession, ServerSession
|
||||
from mcp.shared.context import LifespanContextT, RequestContext
|
||||
from mcp.types import (
|
||||
AudioContent,
|
||||
CreateMessageResult,
|
||||
CreateMessageResultWithTools,
|
||||
ImageContent,
|
||||
ModelPreferences,
|
||||
SamplingMessage,
|
||||
SamplingMessageContentBlock,
|
||||
StopReason,
|
||||
TextContent,
|
||||
ToolChoice,
|
||||
ToolResultContent,
|
||||
ToolUseContent,
|
||||
)
|
||||
from mcp.types import CreateMessageRequestParams as SamplingParams
|
||||
from mcp.types import Tool as MCPTool
|
||||
|
||||
__all__ = ["GoogleGenaiSamplingHandler"]
|
||||
|
||||
|
||||
class GoogleGenaiSamplingHandler:
|
||||
"""Sampling handler that uses the Google GenAI API with tool support.
|
||||
|
||||
Example:
|
||||
```python
|
||||
from google.genai import Client
|
||||
from fastmcp import FastMCP
|
||||
from fastmcp.client.sampling.handlers.google_genai import (
|
||||
GoogleGenaiSamplingHandler,
|
||||
)
|
||||
|
||||
handler = GoogleGenaiSamplingHandler(
|
||||
default_model="gemini-2.0-flash",
|
||||
client=Client(),
|
||||
)
|
||||
|
||||
server = FastMCP(sampling_handler=handler)
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
default_model: str,
|
||||
client: GoogleGenaiClient | None = None,
|
||||
thinking_budget: int | None = None,
|
||||
) -> None:
|
||||
self.client: GoogleGenaiClient = client or GoogleGenaiClient()
|
||||
self.default_model: str = default_model
|
||||
self.thinking_budget: int | None = thinking_budget
|
||||
|
||||
async def __call__(
|
||||
self,
|
||||
messages: list[SamplingMessage],
|
||||
params: SamplingParams,
|
||||
context: RequestContext[ServerSession, LifespanContextT]
|
||||
| RequestContext[ClientSession, LifespanContextT],
|
||||
) -> CreateMessageResult | CreateMessageResultWithTools:
|
||||
contents: list[Content] = _convert_messages_to_google_genai_content(messages)
|
||||
|
||||
# Convert MCP tools to Google GenAI format
|
||||
google_tools: list[GoogleTool] | None = None
|
||||
tool_config: ToolConfig | None = None
|
||||
|
||||
if params.tools:
|
||||
google_tools = [
|
||||
_convert_tool_to_google_genai(tool) for tool in params.tools
|
||||
]
|
||||
tool_config = _convert_tool_choice_to_google_genai(params.toolChoice)
|
||||
|
||||
# Select the model based on preferences
|
||||
selected_model = self._get_model(model_preferences=params.modelPreferences)
|
||||
|
||||
# Configure thinking if a budget is specified
|
||||
thinking_config = (
|
||||
ThinkingConfig(thinking_budget=self.thinking_budget)
|
||||
if self.thinking_budget is not None
|
||||
else None
|
||||
)
|
||||
|
||||
response: GenerateContentResponse = (
|
||||
await self.client.aio.models.generate_content(
|
||||
model=selected_model,
|
||||
contents=contents,
|
||||
config=GenerateContentConfig(
|
||||
system_instruction=params.systemPrompt,
|
||||
temperature=params.temperature,
|
||||
max_output_tokens=params.maxTokens,
|
||||
stop_sequences=params.stopSequences,
|
||||
thinking_config=thinking_config,
|
||||
tools=google_tools, # ty: ignore[invalid-argument-type]
|
||||
tool_config=tool_config,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
# Return appropriate result type based on whether tools were provided
|
||||
if params.tools:
|
||||
return _response_to_result_with_tools(response, selected_model)
|
||||
return _response_to_create_message_result(response, selected_model)
|
||||
|
||||
def _get_model(self, model_preferences: ModelPreferences | None) -> str:
|
||||
if model_preferences and model_preferences.hints:
|
||||
for hint in model_preferences.hints:
|
||||
if hint.name and hint.name.startswith("gemini"):
|
||||
return hint.name
|
||||
return self.default_model
|
||||
|
||||
|
||||
def _convert_tool_to_google_genai(tool: MCPTool) -> GoogleTool:
|
||||
"""Convert an MCP Tool to Google GenAI format.
|
||||
|
||||
Google's parameters_json_schema accepts standard JSON Schema format,
|
||||
so we pass tool.inputSchema directly without conversion.
|
||||
"""
|
||||
return GoogleTool(
|
||||
function_declarations=[
|
||||
FunctionDeclaration(
|
||||
name=tool.name,
|
||||
description=tool.description or "",
|
||||
parameters_json_schema=tool.inputSchema,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def _convert_tool_choice_to_google_genai(tool_choice: ToolChoice | None) -> ToolConfig:
|
||||
"""Convert MCP ToolChoice to Google GenAI ToolConfig."""
|
||||
if tool_choice is None:
|
||||
return ToolConfig(
|
||||
function_calling_config=FunctionCallingConfig(
|
||||
mode=FunctionCallingConfigMode.AUTO
|
||||
)
|
||||
)
|
||||
|
||||
if tool_choice.mode == "required":
|
||||
return ToolConfig(
|
||||
function_calling_config=FunctionCallingConfig(
|
||||
mode=FunctionCallingConfigMode.ANY
|
||||
)
|
||||
)
|
||||
if tool_choice.mode == "none":
|
||||
return ToolConfig(
|
||||
function_calling_config=FunctionCallingConfig(
|
||||
mode=FunctionCallingConfigMode.NONE
|
||||
)
|
||||
)
|
||||
|
||||
# Default to AUTO for "auto" or any other value
|
||||
return ToolConfig(
|
||||
function_calling_config=FunctionCallingConfig(
|
||||
mode=FunctionCallingConfigMode.AUTO
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _sampling_content_to_google_genai_part(
|
||||
content: TextContent
|
||||
| ImageContent
|
||||
| AudioContent
|
||||
| ToolUseContent
|
||||
| ToolResultContent,
|
||||
) -> Part:
|
||||
"""Convert MCP content to Google GenAI Part."""
|
||||
if isinstance(content, TextContent):
|
||||
return Part(text=content.text)
|
||||
|
||||
if isinstance(content, ToolUseContent):
|
||||
# Note: thought_signature bypass is required for manually constructed tool calls.
|
||||
# Google's Gemini 3+ models enforce thought signature validation for function calls.
|
||||
# Since we're constructing these Parts from MCP protocol data (not from model responses),
|
||||
# they lack legitimate signatures. The bypass value allows validation to pass.
|
||||
# See: https://ai.google.dev/gemini-api/docs/thought-signatures
|
||||
return Part(
|
||||
function_call=FunctionCall(
|
||||
name=content.name,
|
||||
args=content.input,
|
||||
),
|
||||
thought_signature=b"skip_thought_signature_validator",
|
||||
)
|
||||
|
||||
if isinstance(content, ToolResultContent):
|
||||
# Extract text from tool result content
|
||||
result_parts: list[str] = []
|
||||
if content.content:
|
||||
for item in content.content:
|
||||
if isinstance(item, TextContent):
|
||||
result_parts.append(item.text)
|
||||
else:
|
||||
msg = f"Unsupported tool result content type: {type(item).__name__}"
|
||||
raise ValueError(msg)
|
||||
result_text = "".join(result_parts)
|
||||
|
||||
# Extract function name from toolUseId
|
||||
# Our IDs are formatted as "{function_name}_{uuid8}", so extract the name.
|
||||
# Note: This is a limitation of MCP's ToolResultContent which only carries
|
||||
# toolUseId, while Google's FunctionResponse requires the function name.
|
||||
tool_use_id = content.toolUseId
|
||||
if "_" in tool_use_id:
|
||||
# Split and rejoin all but the last part (the UUID suffix)
|
||||
parts = tool_use_id.rsplit("_", 1)
|
||||
function_name = parts[0]
|
||||
else:
|
||||
# Fallback: use the full ID as the name
|
||||
function_name = tool_use_id
|
||||
|
||||
return Part(
|
||||
function_response=FunctionResponse(
|
||||
name=function_name,
|
||||
response={"result": result_text},
|
||||
)
|
||||
)
|
||||
|
||||
msg = f"Unsupported content type: {type(content)}"
|
||||
raise ValueError(msg)
|
||||
|
||||
|
||||
def _convert_messages_to_google_genai_content(
|
||||
messages: Sequence[SamplingMessage],
|
||||
) -> list[Content]:
|
||||
"""Convert MCP messages to Google GenAI content."""
|
||||
google_messages: list[Content] = []
|
||||
|
||||
for message in messages:
|
||||
content = message.content
|
||||
|
||||
# Handle list content (tool calls + results)
|
||||
if isinstance(content, list):
|
||||
parts: list[Part] = []
|
||||
for item in content:
|
||||
parts.append(_sampling_content_to_google_genai_part(item))
|
||||
|
||||
if message.role == "user":
|
||||
google_messages.append(UserContent(parts=parts))
|
||||
elif message.role == "assistant":
|
||||
google_messages.append(ModelContent(parts=parts))
|
||||
else:
|
||||
msg = f"Invalid message role: {message.role}"
|
||||
raise ValueError(msg)
|
||||
continue
|
||||
|
||||
# Handle single content item
|
||||
part = _sampling_content_to_google_genai_part(content)
|
||||
|
||||
if message.role == "user":
|
||||
google_messages.append(UserContent(parts=[part]))
|
||||
elif message.role == "assistant":
|
||||
google_messages.append(ModelContent(parts=[part]))
|
||||
else:
|
||||
msg = f"Invalid message role: {message.role}"
|
||||
raise ValueError(msg)
|
||||
|
||||
return google_messages
|
||||
|
||||
|
||||
def _get_candidate_from_response(response: GenerateContentResponse) -> Candidate:
|
||||
"""Extract the first candidate from a response."""
|
||||
if response.candidates and response.candidates[0]:
|
||||
return response.candidates[0]
|
||||
msg = "No candidate in response from completion."
|
||||
raise ValueError(msg)
|
||||
|
||||
|
||||
def _response_to_create_message_result(
|
||||
response: GenerateContentResponse,
|
||||
model: str,
|
||||
) -> CreateMessageResult:
|
||||
"""Convert Google GenAI response to CreateMessageResult (no tools)."""
|
||||
if not (text := response.text):
|
||||
candidate = _get_candidate_from_response(response)
|
||||
msg = f"No content in response: {candidate.finish_reason}"
|
||||
raise ValueError(msg)
|
||||
|
||||
return CreateMessageResult(
|
||||
content=TextContent(type="text", text=text),
|
||||
role="assistant",
|
||||
model=model,
|
||||
)
|
||||
|
||||
|
||||
def _response_to_result_with_tools(
|
||||
response: GenerateContentResponse,
|
||||
model: str,
|
||||
) -> CreateMessageResultWithTools:
|
||||
"""Convert Google GenAI response to CreateMessageResultWithTools."""
|
||||
candidate = _get_candidate_from_response(response)
|
||||
|
||||
# Determine stop reason and check for function calls
|
||||
stop_reason: StopReason
|
||||
finish_reason = candidate.finish_reason
|
||||
has_function_calls = False
|
||||
|
||||
if candidate.content and candidate.content.parts:
|
||||
for part in candidate.content.parts:
|
||||
if part.function_call is not None:
|
||||
has_function_calls = True
|
||||
break
|
||||
|
||||
if has_function_calls:
|
||||
stop_reason = "toolUse"
|
||||
elif finish_reason == "STOP":
|
||||
stop_reason = "endTurn"
|
||||
elif finish_reason == "MAX_TOKENS":
|
||||
stop_reason = "maxTokens"
|
||||
else:
|
||||
stop_reason = "endTurn"
|
||||
|
||||
# Build content list
|
||||
content: list[SamplingMessageContentBlock] = []
|
||||
|
||||
if candidate.content and candidate.content.parts:
|
||||
for part in candidate.content.parts:
|
||||
# Note: Skip thought parts from thinking_config - not relevant for MCP responses
|
||||
if part.text:
|
||||
content.append(TextContent(type="text", text=part.text))
|
||||
elif part.function_call is not None:
|
||||
fc = part.function_call
|
||||
fc_name: str = fc.name or "unknown"
|
||||
content.append(
|
||||
ToolUseContent(
|
||||
type="tool_use",
|
||||
id=f"{fc_name}_{uuid4().hex[:8]}", # Generate unique ID
|
||||
name=fc_name,
|
||||
input=dict(fc.args) if fc.args else {},
|
||||
)
|
||||
)
|
||||
|
||||
if not content:
|
||||
raise ValueError("No content in response from completion")
|
||||
|
||||
return CreateMessageResultWithTools(
|
||||
content=content,
|
||||
role="assistant",
|
||||
model=model,
|
||||
stopReason=stop_reason,
|
||||
)
|
||||
360
tests/client/sampling/handlers/test_google_genai_handler.py
Normal file
360
tests/client/sampling/handlers/test_google_genai_handler.py
Normal file
|
|
@ -0,0 +1,360 @@
|
|||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
try:
|
||||
from google.genai import Client as GoogleGenaiClient
|
||||
from google.genai.types import (
|
||||
Candidate,
|
||||
FunctionCall,
|
||||
FunctionCallingConfigMode,
|
||||
GenerateContentResponse,
|
||||
ModelContent,
|
||||
Part,
|
||||
UserContent,
|
||||
)
|
||||
from mcp.types import (
|
||||
CreateMessageResult,
|
||||
ModelHint,
|
||||
ModelPreferences,
|
||||
TextContent,
|
||||
ToolChoice,
|
||||
ToolResultContent,
|
||||
ToolUseContent,
|
||||
)
|
||||
|
||||
from fastmcp.client.sampling.handlers.google_genai import (
|
||||
GoogleGenaiSamplingHandler,
|
||||
_convert_messages_to_google_genai_content,
|
||||
_convert_tool_choice_to_google_genai,
|
||||
_response_to_create_message_result,
|
||||
_response_to_result_with_tools,
|
||||
_sampling_content_to_google_genai_part,
|
||||
)
|
||||
|
||||
GOOGLE_GENAI_AVAILABLE = True
|
||||
except ImportError:
|
||||
GOOGLE_GENAI_AVAILABLE = False
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not GOOGLE_GENAI_AVAILABLE, reason="google-genai not installed"
|
||||
)
|
||||
|
||||
|
||||
def test_convert_sampling_messages_to_google_genai_content():
|
||||
from mcp.types import SamplingMessage, TextContent
|
||||
|
||||
msgs = _convert_messages_to_google_genai_content(
|
||||
messages=[
|
||||
SamplingMessage(
|
||||
role="user", content=TextContent(type="text", text="hello")
|
||||
),
|
||||
SamplingMessage(
|
||||
role="assistant", content=TextContent(type="text", text="ok")
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
assert len(msgs) == 2
|
||||
assert isinstance(msgs[0], UserContent)
|
||||
assert isinstance(msgs[1], ModelContent)
|
||||
assert msgs[0].parts[0].text == "hello"
|
||||
assert msgs[1].parts[0].text == "ok"
|
||||
|
||||
|
||||
def test_convert_to_google_genai_messages_raises_on_non_text():
|
||||
from mcp.types import SamplingMessage
|
||||
|
||||
from fastmcp.utilities.types import Image
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
_convert_messages_to_google_genai_content(
|
||||
messages=[
|
||||
SamplingMessage(
|
||||
role="user",
|
||||
content=Image(data=b"abc").to_image_content(),
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def test_get_model():
|
||||
mock_client = MagicMock(spec=GoogleGenaiClient)
|
||||
handler = GoogleGenaiSamplingHandler(
|
||||
default_model="fallback-model", client=mock_client
|
||||
)
|
||||
|
||||
# Test with Gemini model hint
|
||||
prefs = ModelPreferences(hints=[ModelHint(name="gemini-2.0-flash-exp")])
|
||||
assert handler._get_model(prefs) == "gemini-2.0-flash-exp"
|
||||
|
||||
# Test with None
|
||||
assert handler._get_model(None) == "fallback-model"
|
||||
|
||||
# Test with empty hints
|
||||
prefs_empty = ModelPreferences(hints=[])
|
||||
assert handler._get_model(prefs_empty) == "fallback-model"
|
||||
|
||||
# Test with non-Gemini hint falls back to default
|
||||
prefs_other = ModelPreferences(hints=[ModelHint(name="gpt-4o")])
|
||||
assert handler._get_model(prefs_other) == "fallback-model"
|
||||
|
||||
# Test with mixed hints selects first Gemini model
|
||||
prefs_mixed = ModelPreferences(
|
||||
hints=[ModelHint(name="claude-3.5-sonnet"), ModelHint(name="gemini-2.0-flash")]
|
||||
)
|
||||
assert handler._get_model(prefs_mixed) == "gemini-2.0-flash"
|
||||
|
||||
|
||||
async def test_response_to_create_message_result():
|
||||
# Create a mock response
|
||||
mock_response = MagicMock(spec=GenerateContentResponse)
|
||||
mock_response.text = "HELPFUL CONTENT FROM GEMINI"
|
||||
|
||||
result: CreateMessageResult = _response_to_create_message_result(
|
||||
response=mock_response, model="gemini-2.0-flash-exp"
|
||||
)
|
||||
assert result == CreateMessageResult(
|
||||
content=TextContent(type="text", text="HELPFUL CONTENT FROM GEMINI"),
|
||||
role="assistant",
|
||||
model="gemini-2.0-flash-exp",
|
||||
)
|
||||
|
||||
|
||||
def test_convert_tool_choice_to_google_genai():
|
||||
# Test auto mode
|
||||
result = _convert_tool_choice_to_google_genai(ToolChoice(mode="auto"))
|
||||
assert result.function_calling_config is not None
|
||||
assert result.function_calling_config.mode == FunctionCallingConfigMode.AUTO
|
||||
|
||||
# Test required mode
|
||||
result = _convert_tool_choice_to_google_genai(ToolChoice(mode="required"))
|
||||
assert result.function_calling_config is not None
|
||||
assert result.function_calling_config.mode == FunctionCallingConfigMode.ANY
|
||||
|
||||
# Test none mode
|
||||
result = _convert_tool_choice_to_google_genai(ToolChoice(mode="none"))
|
||||
assert result.function_calling_config is not None
|
||||
assert result.function_calling_config.mode == FunctionCallingConfigMode.NONE
|
||||
|
||||
# Test None (defaults to auto)
|
||||
result = _convert_tool_choice_to_google_genai(None)
|
||||
assert result.function_calling_config is not None
|
||||
assert result.function_calling_config.mode == FunctionCallingConfigMode.AUTO
|
||||
|
||||
|
||||
def test_sampling_content_to_google_genai_part_tool_use():
|
||||
"""Test converting ToolUseContent to Google GenAI Part with FunctionCall."""
|
||||
content = ToolUseContent(
|
||||
type="tool_use",
|
||||
id="get_weather_abc123",
|
||||
name="get_weather",
|
||||
input={"city": "London"},
|
||||
)
|
||||
|
||||
part = _sampling_content_to_google_genai_part(content)
|
||||
|
||||
assert part.function_call is not None
|
||||
assert part.function_call.name == "get_weather"
|
||||
assert part.function_call.args == {"city": "London"}
|
||||
|
||||
|
||||
def test_sampling_content_to_google_genai_part_tool_result():
|
||||
"""Test converting ToolResultContent to Google GenAI Part with FunctionResponse."""
|
||||
content = ToolResultContent(
|
||||
type="tool_result",
|
||||
toolUseId="get_weather_abc123",
|
||||
content=[TextContent(type="text", text="Weather is sunny")],
|
||||
)
|
||||
|
||||
part = _sampling_content_to_google_genai_part(content)
|
||||
|
||||
assert part.function_response is not None
|
||||
# Function name is extracted from toolUseId by removing the UUID suffix
|
||||
assert part.function_response.name == "get_weather"
|
||||
assert part.function_response.response == {"result": "Weather is sunny"}
|
||||
|
||||
|
||||
def test_sampling_content_to_google_genai_part_tool_result_empty():
|
||||
"""Test converting empty ToolResultContent to Google GenAI Part."""
|
||||
content = ToolResultContent(
|
||||
type="tool_result",
|
||||
toolUseId="my_tool_xyz789",
|
||||
content=[],
|
||||
)
|
||||
|
||||
part = _sampling_content_to_google_genai_part(content)
|
||||
|
||||
assert part.function_response is not None
|
||||
assert part.function_response.name == "my_tool"
|
||||
assert part.function_response.response == {"result": ""}
|
||||
|
||||
|
||||
def test_sampling_content_to_google_genai_part_tool_result_no_underscore():
|
||||
"""Test ToolResultContent when toolUseId has no underscore (fallback)."""
|
||||
content = ToolResultContent(
|
||||
type="tool_result",
|
||||
toolUseId="simplefunction",
|
||||
content=[TextContent(type="text", text="Result")],
|
||||
)
|
||||
|
||||
part = _sampling_content_to_google_genai_part(content)
|
||||
|
||||
# When no underscore, the full ID is used as the name
|
||||
assert part.function_response is not None
|
||||
assert part.function_response.name == "simplefunction"
|
||||
|
||||
|
||||
def test_convert_messages_with_tool_use():
|
||||
"""Test converting messages containing ToolUseContent."""
|
||||
from mcp.types import SamplingMessage
|
||||
|
||||
msgs = _convert_messages_to_google_genai_content(
|
||||
messages=[
|
||||
SamplingMessage(
|
||||
role="user",
|
||||
content=TextContent(type="text", text="What's the weather?"),
|
||||
),
|
||||
SamplingMessage(
|
||||
role="assistant",
|
||||
content=ToolUseContent(
|
||||
type="tool_use",
|
||||
id="get_weather_123",
|
||||
name="get_weather",
|
||||
input={"city": "NYC"},
|
||||
),
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
assert len(msgs) == 2
|
||||
assert isinstance(msgs[0], UserContent)
|
||||
assert isinstance(msgs[1], ModelContent)
|
||||
assert msgs[1].parts[0].function_call is not None
|
||||
assert msgs[1].parts[0].function_call.name == "get_weather"
|
||||
|
||||
|
||||
def test_convert_messages_with_tool_result():
|
||||
"""Test converting messages containing ToolResultContent."""
|
||||
from mcp.types import SamplingMessage
|
||||
|
||||
msgs = _convert_messages_to_google_genai_content(
|
||||
messages=[
|
||||
SamplingMessage(
|
||||
role="user",
|
||||
content=ToolResultContent(
|
||||
type="tool_result",
|
||||
toolUseId="get_weather_123",
|
||||
content=[TextContent(type="text", text="Sunny, 72°F")],
|
||||
),
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
assert len(msgs) == 1
|
||||
assert isinstance(msgs[0], UserContent)
|
||||
assert msgs[0].parts[0].function_response is not None
|
||||
assert msgs[0].parts[0].function_response.name == "get_weather"
|
||||
|
||||
|
||||
def test_convert_messages_with_multiple_content_blocks():
|
||||
"""Test converting messages with multiple content blocks (list content)."""
|
||||
from mcp.types import SamplingMessage
|
||||
|
||||
msgs = _convert_messages_to_google_genai_content(
|
||||
messages=[
|
||||
SamplingMessage(
|
||||
role="user",
|
||||
content=[
|
||||
TextContent(type="text", text="I need weather info."),
|
||||
ToolResultContent(
|
||||
type="tool_result",
|
||||
toolUseId="get_weather_xyz",
|
||||
content=[TextContent(type="text", text="Cloudy")],
|
||||
),
|
||||
],
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
assert len(msgs) == 1
|
||||
assert isinstance(msgs[0], UserContent)
|
||||
assert len(msgs[0].parts) == 2
|
||||
assert msgs[0].parts[0].text == "I need weather info."
|
||||
assert msgs[0].parts[1].function_response is not None
|
||||
|
||||
|
||||
def test_response_to_result_with_tools_text_only():
|
||||
"""Test _response_to_result_with_tools with a text-only response."""
|
||||
mock_candidate = MagicMock(spec=Candidate)
|
||||
mock_candidate.content = MagicMock()
|
||||
mock_candidate.content.parts = [Part(text="Here's the answer")]
|
||||
mock_candidate.finish_reason = "STOP"
|
||||
|
||||
mock_response = MagicMock(spec=GenerateContentResponse)
|
||||
mock_response.candidates = [mock_candidate]
|
||||
|
||||
result = _response_to_result_with_tools(mock_response, model="gemini-2.0-flash")
|
||||
|
||||
assert result.role == "assistant"
|
||||
assert result.model == "gemini-2.0-flash"
|
||||
assert result.stopReason == "endTurn"
|
||||
assert isinstance(result.content, list)
|
||||
assert len(result.content) == 1
|
||||
assert result.content[0].type == "text"
|
||||
assert isinstance(result.content[0], TextContent)
|
||||
assert result.content[0].text == "Here's the answer"
|
||||
|
||||
|
||||
def test_response_to_result_with_tools_function_call():
|
||||
"""Test _response_to_result_with_tools with a function call response."""
|
||||
mock_candidate = MagicMock(spec=Candidate)
|
||||
mock_candidate.content = MagicMock()
|
||||
mock_candidate.content.parts = [
|
||||
Part(function_call=FunctionCall(name="get_weather", args={"city": "Paris"}))
|
||||
]
|
||||
mock_candidate.finish_reason = "STOP"
|
||||
|
||||
mock_response = MagicMock(spec=GenerateContentResponse)
|
||||
mock_response.candidates = [mock_candidate]
|
||||
|
||||
result = _response_to_result_with_tools(mock_response, model="gemini-2.0-flash")
|
||||
|
||||
assert result.stopReason == "toolUse"
|
||||
assert isinstance(result.content, list)
|
||||
assert len(result.content) == 1
|
||||
tool_use = result.content[0]
|
||||
assert isinstance(tool_use, ToolUseContent)
|
||||
assert tool_use.type == "tool_use"
|
||||
assert tool_use.name == "get_weather"
|
||||
assert tool_use.input == {"city": "Paris"}
|
||||
# ID should be in format "get_weather_{uuid}"
|
||||
assert tool_use.id.startswith("get_weather_")
|
||||
|
||||
|
||||
def test_response_to_result_with_tools_mixed_content():
|
||||
"""Test _response_to_result_with_tools with text and function call."""
|
||||
mock_candidate = MagicMock(spec=Candidate)
|
||||
mock_candidate.content = MagicMock()
|
||||
mock_candidate.content.parts = [
|
||||
Part(text="Let me check that for you."),
|
||||
Part(function_call=FunctionCall(name="search", args={"query": "test"})),
|
||||
]
|
||||
mock_candidate.finish_reason = "STOP"
|
||||
|
||||
mock_response = MagicMock(spec=GenerateContentResponse)
|
||||
mock_response.candidates = [mock_candidate]
|
||||
|
||||
result = _response_to_result_with_tools(mock_response, model="gemini-2.0-flash")
|
||||
|
||||
assert result.stopReason == "toolUse"
|
||||
assert isinstance(result.content, list)
|
||||
assert len(result.content) == 2
|
||||
text_content = result.content[0]
|
||||
assert isinstance(text_content, TextContent)
|
||||
assert text_content.type == "text"
|
||||
assert text_content.text == "Let me check that for you."
|
||||
tool_use = result.content[1]
|
||||
assert isinstance(tool_use, ToolUseContent)
|
||||
assert tool_use.type == "tool_use"
|
||||
assert tool_use.name == "search"
|
||||
92
uv.lock
generated
92
uv.lock
generated
|
|
@ -796,6 +796,9 @@ azure = [
|
|||
code-mode = [
|
||||
{ name = "pydantic-monty" },
|
||||
]
|
||||
gemini = [
|
||||
{ name = "google-genai" },
|
||||
]
|
||||
openai = [
|
||||
{ name = "openai" },
|
||||
]
|
||||
|
|
@ -807,7 +810,7 @@ tasks = [
|
|||
dev = [
|
||||
{ name = "dirty-equals" },
|
||||
{ name = "fastapi" },
|
||||
{ name = "fastmcp", extra = ["anthropic", "apps", "azure", "code-mode", "openai", "tasks"] },
|
||||
{ name = "fastmcp", extra = ["anthropic", "apps", "azure", "code-mode", "gemini", "openai", "tasks"] },
|
||||
{ name = "inline-snapshot", extra = ["dirty-equals"] },
|
||||
{ name = "ipython", version = "8.38.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
|
||||
{ name = "ipython", version = "9.10.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" },
|
||||
|
|
@ -840,6 +843,7 @@ requires-dist = [
|
|||
{ name = "azure-identity", marker = "extra == 'azure'", specifier = ">=1.16.0" },
|
||||
{ name = "cyclopts", specifier = ">=4.0.0" },
|
||||
{ name = "exceptiongroup", specifier = ">=1.2.2" },
|
||||
{ name = "google-genai", marker = "extra == 'gemini'", specifier = ">=1.18.0" },
|
||||
{ name = "httpx", specifier = ">=0.28.1,<1.0" },
|
||||
{ name = "jsonref", specifier = ">=1.1.0" },
|
||||
{ name = "jsonschema-path", specifier = ">=0.3.4" },
|
||||
|
|
@ -863,13 +867,13 @@ requires-dist = [
|
|||
{ name = "watchfiles", specifier = ">=1.0.0" },
|
||||
{ name = "websockets", specifier = ">=15.0.1" },
|
||||
]
|
||||
provides-extras = ["anthropic", "apps", "azure", "code-mode", "openai", "tasks"]
|
||||
provides-extras = ["anthropic", "apps", "azure", "code-mode", "gemini", "openai", "tasks"]
|
||||
|
||||
[package.metadata.requires-dev]
|
||||
dev = [
|
||||
{ name = "dirty-equals", specifier = ">=0.9.0" },
|
||||
{ name = "fastapi", specifier = ">=0.115.12" },
|
||||
{ name = "fastmcp", extras = ["anthropic", "apps", "azure", "code-mode", "openai", "tasks"] },
|
||||
{ name = "fastmcp", extras = ["anthropic", "apps", "azure", "code-mode", "gemini", "openai", "tasks"] },
|
||||
{ name = "inline-snapshot", extras = ["dirty-equals"], specifier = ">=0.27.2" },
|
||||
{ name = "ipython", specifier = ">=8.12.3" },
|
||||
{ name = "loq", specifier = ">=0.1.0a3" },
|
||||
|
|
@ -894,6 +898,46 @@ dev = [
|
|||
{ name = "ty", specifier = ">=0.0.15" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "google-auth"
|
||||
version = "2.48.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "cryptography" },
|
||||
{ name = "pyasn1-modules" },
|
||||
{ name = "rsa" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/0c/41/242044323fbd746615884b1c16639749e73665b718209946ebad7ba8a813/google_auth-2.48.0.tar.gz", hash = "sha256:4f7e706b0cd3208a3d940a19a822c37a476ddba5450156c3e6624a71f7c841ce", size = 326522, upload-time = "2026-01-26T19:22:47.157Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/83/1d/d6466de3a5249d35e832a52834115ca9d1d0de6abc22065f049707516d47/google_auth-2.48.0-py3-none-any.whl", hash = "sha256:2e2a537873d449434252a9632c28bfc268b0adb1e53f9fb62afc5333a975903f", size = 236499, upload-time = "2026-01-26T19:22:45.099Z" },
|
||||
]
|
||||
|
||||
[package.optional-dependencies]
|
||||
requests = [
|
||||
{ name = "requests" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "google-genai"
|
||||
version = "1.65.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "anyio" },
|
||||
{ name = "distro" },
|
||||
{ name = "google-auth", extra = ["requests"] },
|
||||
{ name = "httpx" },
|
||||
{ name = "pydantic" },
|
||||
{ name = "requests" },
|
||||
{ name = "sniffio" },
|
||||
{ name = "tenacity" },
|
||||
{ name = "typing-extensions" },
|
||||
{ name = "websockets" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/79/f9/cc1191c2540d6a4e24609a586c4ed45d2db57cfef47931c139ee70e5874a/google_genai-1.65.0.tar.gz", hash = "sha256:d470eb600af802d58a79c7f13342d9ea0d05d965007cae8f76c7adff3d7a4750", size = 497206, upload-time = "2026-02-26T00:20:33.824Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/68/3c/3fea4e7c91357c71782d7dcaad7a2577d636c90317e003386893c25bc62c/google_genai-1.65.0-py3-none-any.whl", hash = "sha256:68c025205856919bc03edb0155c11b4b833810b7ce17ad4b7a9eeba5158f6c44", size = 724429, upload-time = "2026-02-26T00:20:32.186Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "googleapis-common-protos"
|
||||
version = "1.72.0"
|
||||
|
|
@ -1864,6 +1908,27 @@ redis = [
|
|||
{ name = "redis" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pyasn1"
|
||||
version = "0.6.2"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/fe/b6/6e630dff89739fcd427e3f72b3d905ce0acb85a45d4ec3e2678718a3487f/pyasn1-0.6.2.tar.gz", hash = "sha256:9b59a2b25ba7e4f8197db7686c09fb33e658b98339fadb826e9512629017833b", size = 146586, upload-time = "2026-01-16T18:04:18.534Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/44/b5/a96872e5184f354da9c84ae119971a0a4c221fe9b27a4d94bd43f2596727/pyasn1-0.6.2-py3-none-any.whl", hash = "sha256:1eb26d860996a18e9b6ed05e7aae0e9fc21619fcee6af91cca9bad4fbea224bf", size = 83371, upload-time = "2026-01-16T18:04:17.174Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pyasn1-modules"
|
||||
version = "0.4.2"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "pyasn1" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/e9/e6/78ebbb10a8c8e4b61a59249394a4a594c1a7af95593dc933a349c8d00964/pyasn1_modules-0.4.2.tar.gz", hash = "sha256:677091de870a80aae844b1ca6134f54652fa2c8c5a52aa396440ac3106e941e6", size = 307892, upload-time = "2025-03-28T02:41:22.17Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/47/8d/d529b5d697919ba8c11ad626e835d4039be708a35b0d22de83a269a6682c/pyasn1_modules-0.4.2-py3-none-any.whl", hash = "sha256:29253a9207ce32b64c3ac6600edc75368f98473906e8fd1043bd6b5b1de2c14a", size = 181259, upload-time = "2025-03-28T02:41:19.028Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pycparser"
|
||||
version = "3.0"
|
||||
|
|
@ -2688,6 +2753,18 @@ wheels = [
|
|||
{ url = "https://files.pythonhosted.org/packages/d1/b7/b95708304cd49b7b6f82fdd039f1748b66ec2b21d6a45180910802f1abf1/rpds_py-0.30.0-pp311-pypy311_pp73-musllinux_1_2_x86_64.whl", hash = "sha256:ac37f9f516c51e5753f27dfdef11a88330f04de2d564be3991384b2f3535d02e", size = 562191, upload-time = "2025-11-30T20:24:36.853Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rsa"
|
||||
version = "4.9.1"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "pyasn1" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/da/8a/22b7beea3ee0d44b1916c0c1cb0ee3af23b700b6da9f04991899d0c555d4/rsa-4.9.1.tar.gz", hash = "sha256:e7bdbfdb5497da4c07dfd35530e1a902659db6ff241e39d9953cad06ebd0ae75", size = 29034, upload-time = "2025-04-16T09:51:18.218Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/64/8d/0133e4eb4beed9e425d9a98ed6e081a55d195481b7632472be1af08d2f6b/rsa-4.9.1-py3-none-any.whl", hash = "sha256:68635866661c6836b8d39430f97a996acbd61bfa49406748ea243539fe239762", size = 34696, upload-time = "2025-04-16T09:51:17.142Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ruff"
|
||||
version = "0.15.1"
|
||||
|
|
@ -2815,6 +2892,15 @@ wheels = [
|
|||
{ url = "https://files.pythonhosted.org/packages/d1/b1/74babcc824a57904e919f3af16d86c08b524c0691504baf038ef2d7f655c/taskgroup-0.2.2-py2.py3-none-any.whl", hash = "sha256:e2c53121609f4ae97303e9ea1524304b4de6faf9eb2c9280c7f87976479a52fb", size = 14237, upload-time = "2025-01-03T09:24:11.41Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tenacity"
|
||||
version = "9.1.4"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/47/c6/ee486fd809e357697ee8a44d3d69222b344920433d3b6666ccd9b374630c/tenacity-9.1.4.tar.gz", hash = "sha256:adb31d4c263f2bd041081ab33b498309a57c77f9acf2db65aadf0898179cf93a", size = 49413, upload-time = "2026-02-07T10:45:33.841Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/d7/c1/eb8f9debc45d3b7918a32ab756658a0904732f75e555402972246b0b8e71/tenacity-9.1.4-py3-none-any.whl", hash = "sha256:6095a360c919085f28c6527de529e76a06ad89b23659fa881ae0649b867a9d55", size = 28926, upload-time = "2026-02-07T10:45:32.24Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tomli"
|
||||
version = "2.4.0"
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue