You've already forked adk-python
mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
Compare commits
1 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 8b0581f5c5 |
@@ -1,70 +0,0 @@
|
||||
# AgentTool with MCP Demo (SSE Mode)
|
||||
|
||||
This demo shows how `AgentTool` works with MCP (Model Context Protocol) toolsets using **SSE mode**.
|
||||
|
||||
## SSE vs Stdio Mode
|
||||
|
||||
This demo uses **SSE (Server-Sent Events) mode** where the MCP server runs as a separate HTTP server:
|
||||
- **Remote connection** - Connects to server via HTTP
|
||||
- **Separate process** - Server must be started manually
|
||||
- **Network communication** - Uses HTTP/SSE for messaging
|
||||
|
||||
For the **stdio (subprocess) version**, see [mcp_in_agent_tool_stdio](../mcp_in_agent_tool_stdio/).
|
||||
|
||||
## Setup
|
||||
|
||||
**Start the MCP simple-tool server in SSE mode** (in a separate terminal):
|
||||
```bash
|
||||
# Run the server using uvx (no installation needed)
|
||||
# Port 3000 avoids conflict with adk web (which uses 8000)
|
||||
uvx --from 'git+https://github.com/modelcontextprotocol/python-sdk.git#subdirectory=examples/servers/simple-tool' \
|
||||
mcp-simple-tool --transport sse --port 3000
|
||||
```
|
||||
|
||||
The server should be accessible at `http://localhost:3000/sse`.
|
||||
|
||||
## Running the Demo
|
||||
|
||||
```bash
|
||||
adk web contributing/samples
|
||||
```
|
||||
|
||||
Then select **mcp_in_agent_tool_remote** from the list and interact with the agent.
|
||||
|
||||
## Try These Prompts
|
||||
|
||||
This demo uses **Gemini 2.5 Flash** as the model. Try these prompts:
|
||||
|
||||
1. **Check available tools:**
|
||||
```
|
||||
What tools do you have access to?
|
||||
```
|
||||
|
||||
2. **Fetch and summarize JSON Schema specification:**
|
||||
```
|
||||
Use the mcp_helper to fetch https://json-schema.org/specification and summarize the key features of JSON Schema
|
||||
```
|
||||
|
||||
## Architecture
|
||||
|
||||
```
|
||||
main_agent (root_agent)
|
||||
│
|
||||
└── AgentTool wrapping:
|
||||
│
|
||||
└── mcp_helper (sub_agent)
|
||||
│
|
||||
└── McpToolset (SSE connection)
|
||||
│
|
||||
└── http://localhost:3000/sse
|
||||
│
|
||||
└── MCP simple-tool server
|
||||
│
|
||||
└── Website Fetcher Tool
|
||||
```
|
||||
|
||||
## Related
|
||||
|
||||
- **Issue:** [#1112 - Using agent as tool outside of adk web doesn't exit cleanly](https://github.com/google/adk-python/issues/1112)
|
||||
- **Related Issue:** [#929 - LiteLLM giving error with OpenAI models and Grafana's MCP server](https://github.com/google/adk-python/issues/929)
|
||||
- **Stdio Version:** [mcp_in_agent_tool_stdio](../mcp_in_agent_tool_stdio/) - Uses local subprocess connection
|
||||
@@ -1,15 +0,0 @@
|
||||
# Copyright 2025 Google LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from . import agent
|
||||
@@ -1,70 +0,0 @@
|
||||
# Copyright 2025 Google LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from google.adk.agents import Agent
|
||||
from google.adk.tools import AgentTool
|
||||
from google.adk.tools.mcp_tool import McpToolset
|
||||
from google.adk.tools.mcp_tool.mcp_session_manager import SseConnectionParams
|
||||
|
||||
# Create MCP toolset
|
||||
# This uses the simple-tool MCP server via SSE
|
||||
# You need to start the MCP server separately (see README.md)
|
||||
mcp_toolset = McpToolset(
|
||||
connection_params=SseConnectionParams(
|
||||
url="http://localhost:3000/sse",
|
||||
timeout=10.0,
|
||||
sse_read_timeout=300.0,
|
||||
)
|
||||
)
|
||||
|
||||
# Create sub-agent with MCP tools
|
||||
# This agent has direct access to MCP tools
|
||||
sub_agent = Agent(
|
||||
name="mcp_helper",
|
||||
model="gemini-2.5-flash",
|
||||
description=(
|
||||
"A helpful assistant with access to MCP tools for fetching websites."
|
||||
),
|
||||
instruction="""You are a helpful assistant with access to MCP tools.
|
||||
|
||||
When the user asks for help:
|
||||
1. Explain what tools you have available (website fetching)
|
||||
2. Use the appropriate tool if needed
|
||||
3. Provide clear and helpful responses
|
||||
|
||||
You have access to a website fetcher tool via MCP. Use it to fetch and return website content.""",
|
||||
tools=[mcp_toolset],
|
||||
)
|
||||
|
||||
# Wrap sub-agent as an AgentTool
|
||||
# This allows the main agent to delegate tasks to the sub-agent
|
||||
# The sub-agent has access to MCP tools for fetching websites
|
||||
mcp_agent_tool = AgentTool(agent=sub_agent)
|
||||
|
||||
# Create main agent
|
||||
# This agent can delegate to the sub-agent via AgentTool
|
||||
root_agent = Agent(
|
||||
name="main_agent",
|
||||
model="gemini-2.5-flash",
|
||||
description="Main agent that can delegate to a sub-agent with MCP tools.",
|
||||
instruction="""You are a helpful assistant. You have access to a sub-agent (mcp_helper)
|
||||
that has MCP tools for fetching websites.
|
||||
|
||||
When the user asks for help:
|
||||
- If they need to fetch a website, call the mcp_helper tool
|
||||
- Otherwise, respond directly
|
||||
|
||||
Always be helpful and explain what you're doing.""",
|
||||
tools=[mcp_agent_tool],
|
||||
)
|
||||
@@ -1,70 +0,0 @@
|
||||
# AgentTool with MCP Demo (Stdio Mode)
|
||||
|
||||
This demo shows how `AgentTool` works with MCP (Model Context Protocol) toolsets using **stdio mode**.
|
||||
|
||||
## Stdio vs SSE Mode
|
||||
|
||||
This demo uses **stdio mode** where the MCP server runs as a subprocess:
|
||||
- **Simpler setup** - No need to start a separate server
|
||||
- **Auto-launched** - Server starts automatically when agent runs
|
||||
- **Local process** - Uses stdin/stdout for communication
|
||||
|
||||
For the **SSE (remote server) version**, see [mcp_in_agent_tool_remote](../mcp_in_agent_tool_remote/).
|
||||
|
||||
## Setup
|
||||
|
||||
**No installation required!** The MCP server will be launched automatically using `uvx` when you run the agent.
|
||||
|
||||
The demo uses `uvx` to fetch and run the MCP simple-tool server directly from the GitHub repository's subdirectory:
|
||||
```bash
|
||||
uvx --from 'git+https://github.com/modelcontextprotocol/python-sdk.git#subdirectory=examples/servers/simple-tool' \
|
||||
mcp-simple-tool
|
||||
```
|
||||
|
||||
This happens automatically via the stdio connection when the agent starts.
|
||||
|
||||
## Running the Demo
|
||||
|
||||
```bash
|
||||
adk web contributing/samples
|
||||
```
|
||||
|
||||
Then select **mcp_in_agent_tool_stdio** from the list and interact with the agent.
|
||||
|
||||
## Try These Prompts
|
||||
|
||||
This demo uses **Gemini 2.5 Flash** as the model. Try these prompts:
|
||||
|
||||
1. **Check available tools:**
|
||||
```
|
||||
What tools do you have access to?
|
||||
```
|
||||
|
||||
2. **Fetch and summarize JSON Schema specification:**
|
||||
```
|
||||
Use the mcp_helper to fetch https://json-schema.org/specification and summarize the key features of JSON Schema
|
||||
```
|
||||
|
||||
## Architecture
|
||||
|
||||
```
|
||||
main_agent (root_agent)
|
||||
│
|
||||
└── AgentTool wrapping:
|
||||
│
|
||||
└── mcp_helper (sub_agent)
|
||||
│
|
||||
└── McpToolset (stdio connection)
|
||||
│
|
||||
└── MCP Server (subprocess via uvx)
|
||||
│
|
||||
└── uvx --from git+...#subdirectory=... mcp-simple-tool
|
||||
│
|
||||
└── Website Fetcher Tool
|
||||
```
|
||||
|
||||
## Related
|
||||
|
||||
- **Issue:** [#1112 - Using agent as tool outside of adk web doesn't exit cleanly](https://github.com/google/adk-python/issues/1112)
|
||||
- **Related Issue:** [#929 - LiteLLM giving error with OpenAI models and Grafana's MCP server](https://github.com/google/adk-python/issues/929)
|
||||
- **SSE Version:** [mcp_in_agent_tool_remote](../mcp_in_agent_tool_remote/) - Uses remote server connection
|
||||
@@ -1,15 +0,0 @@
|
||||
# Copyright 2025 Google LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from . import agent
|
||||
@@ -1,77 +0,0 @@
|
||||
# Copyright 2025 Google LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from google.adk.agents import Agent
|
||||
from google.adk.tools import AgentTool
|
||||
from google.adk.tools.mcp_tool import McpToolset
|
||||
from google.adk.tools.mcp_tool.mcp_session_manager import StdioConnectionParams
|
||||
from mcp import StdioServerParameters
|
||||
|
||||
# Create MCP toolset
|
||||
# This uses the simple-tool MCP server via stdio
|
||||
# The server will be launched automatically using uvx from the subdirectory
|
||||
mcp_toolset = McpToolset(
|
||||
connection_params=StdioConnectionParams(
|
||||
server_params=StdioServerParameters(
|
||||
command="uvx",
|
||||
args=[
|
||||
"--from",
|
||||
"git+https://github.com/modelcontextprotocol/python-sdk.git#subdirectory=examples/servers/simple-tool",
|
||||
"mcp-simple-tool",
|
||||
],
|
||||
),
|
||||
timeout=10.0,
|
||||
)
|
||||
)
|
||||
|
||||
# Create sub-agent with MCP tools
|
||||
# This agent has direct access to MCP tools
|
||||
sub_agent = Agent(
|
||||
name="mcp_helper",
|
||||
model="gemini-2.5-flash",
|
||||
description=(
|
||||
"A helpful assistant with access to MCP tools for fetching websites."
|
||||
),
|
||||
instruction="""You are a helpful assistant with access to MCP tools.
|
||||
|
||||
When the user asks for help:
|
||||
1. Explain what tools you have available (website fetching)
|
||||
2. Use the appropriate tool if needed
|
||||
3. Provide clear and helpful responses
|
||||
|
||||
You have access to a website fetcher tool via MCP. Use it to fetch and return website content.""",
|
||||
tools=[mcp_toolset],
|
||||
)
|
||||
|
||||
# Wrap sub-agent as an AgentTool
|
||||
# This allows the main agent to delegate tasks to the sub-agent
|
||||
# The sub-agent has access to MCP tools for fetching websites
|
||||
mcp_agent_tool = AgentTool(agent=sub_agent)
|
||||
|
||||
# Create main agent
|
||||
# This agent can delegate to the sub-agent via AgentTool
|
||||
root_agent = Agent(
|
||||
name="main_agent",
|
||||
model="gemini-2.5-flash",
|
||||
description="Main agent that can delegate to a sub-agent with MCP tools.",
|
||||
instruction="""You are a helpful assistant. You have access to a sub-agent (mcp_helper)
|
||||
that has MCP tools for fetching websites.
|
||||
|
||||
When the user asks for help:
|
||||
- If they need to fetch a website, call the mcp_helper tool
|
||||
- Otherwise, respond directly
|
||||
|
||||
Always be helpful and explain what you're doing.""",
|
||||
tools=[mcp_agent_tool],
|
||||
)
|
||||
@@ -15,11 +15,13 @@
|
||||
from .base_plugin import BasePlugin
|
||||
from .logging_plugin import LoggingPlugin
|
||||
from .plugin_manager import PluginManager
|
||||
from .rate_limit_plugin import RateLimitPlugin
|
||||
from .reflect_retry_tool_plugin import ReflectAndRetryToolPlugin
|
||||
|
||||
__all__ = [
|
||||
'BasePlugin',
|
||||
'LoggingPlugin',
|
||||
'PluginManager',
|
||||
'RateLimitPlugin',
|
||||
'ReflectAndRetryToolPlugin',
|
||||
]
|
||||
|
||||
@@ -0,0 +1,163 @@
|
||||
# Copyright 2025 Google LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from typing import Optional
|
||||
|
||||
from ..agents.callback_context import CallbackContext
|
||||
from ..models.llm_request import LlmRequest
|
||||
from ..models.llm_response import LlmResponse
|
||||
from .base_plugin import BasePlugin
|
||||
|
||||
|
||||
class RateLimitPlugin(BasePlugin):
|
||||
"""Plugin that enforces global rate limiting on LLM requests.
|
||||
|
||||
This plugin implements a sliding window rate limiter that restricts the
|
||||
total number of LLM requests across all models within a one-minute window.
|
||||
When the rate limit is exceeded, the plugin blocks (waits) until a slot
|
||||
becomes available.
|
||||
|
||||
Example:
|
||||
```python
|
||||
from google.adk import Agent, Runner
|
||||
from google.adk.plugins.rate_limit_plugin import RateLimitPlugin
|
||||
|
||||
agent = Agent(name="assistant", model="gemini-2.5-flash", ...)
|
||||
|
||||
runner = Runner(
|
||||
agents=[agent],
|
||||
plugins=[
|
||||
RateLimitPlugin(max_requests_per_minute=15)
|
||||
]
|
||||
)
|
||||
```
|
||||
|
||||
Attributes:
|
||||
max_requests_per_minute: Maximum number of requests allowed per minute
|
||||
globally across all models.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
max_requests_per_minute: int = 15,
|
||||
name: str = 'rate_limit_plugin',
|
||||
):
|
||||
"""Initialize the rate limit plugin.
|
||||
|
||||
Args:
|
||||
max_requests_per_minute: Maximum requests allowed per minute globally.
|
||||
name: Name of the plugin instance.
|
||||
"""
|
||||
super().__init__(name)
|
||||
self.max_requests = max_requests_per_minute
|
||||
|
||||
# Track request timestamps globally (all models)
|
||||
# List of timestamps (in seconds since epoch)
|
||||
self._request_timestamps: list[float] = []
|
||||
|
||||
# Lock for thread-safe access to timestamps
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
def _clean_old_timestamps(
|
||||
self, timestamps: list[float], current_time: float
|
||||
) -> list[float]:
|
||||
"""Remove timestamps older than 1 minute from the tracking list.
|
||||
|
||||
Args:
|
||||
timestamps: List of request timestamps.
|
||||
current_time: Current time in seconds since epoch.
|
||||
|
||||
Returns:
|
||||
Filtered list containing only timestamps from the last minute.
|
||||
"""
|
||||
# Keep only timestamps within the last 60 seconds
|
||||
cutoff_time = current_time - 60.0
|
||||
return [ts for ts in timestamps if ts > cutoff_time]
|
||||
|
||||
async def _wait_for_rate_limit(self, current_time: float) -> None:
|
||||
"""Wait until a request slot becomes available.
|
||||
|
||||
Args:
|
||||
current_time: Current time in seconds since epoch.
|
||||
"""
|
||||
while True:
|
||||
async with self._lock:
|
||||
timestamps = self._clean_old_timestamps(
|
||||
self._request_timestamps, time.time()
|
||||
)
|
||||
self._request_timestamps = timestamps
|
||||
|
||||
if len(timestamps) < self.max_requests:
|
||||
# Slot available, exit loop
|
||||
return
|
||||
|
||||
# Calculate wait time until the oldest request falls outside the window
|
||||
oldest_timestamp = timestamps[0]
|
||||
wait_seconds = 60.0 - (time.time() - oldest_timestamp) + 0.1
|
||||
|
||||
# Wait outside the lock to allow other operations
|
||||
if wait_seconds > 0:
|
||||
await asyncio.sleep(wait_seconds)
|
||||
else:
|
||||
# Re-check immediately
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
async def before_model_callback(
|
||||
self, *, callback_context: CallbackContext, llm_request: LlmRequest
|
||||
) -> Optional[LlmResponse]:
|
||||
"""Check and enforce rate limits before each LLM request.
|
||||
|
||||
This callback is invoked before every LLM request. It checks whether
|
||||
the request would exceed the configured global rate limit across all models.
|
||||
If so, it blocks (waits) until the rate limit allows the request.
|
||||
|
||||
Args:
|
||||
callback_context: Context containing agent, user, and session information.
|
||||
llm_request: The LLM request that is about to be sent.
|
||||
|
||||
Returns:
|
||||
None to allow the request to proceed (after waiting if necessary).
|
||||
"""
|
||||
current_time = time.time()
|
||||
|
||||
async with self._lock:
|
||||
# Clean old timestamps
|
||||
timestamps = self._clean_old_timestamps(
|
||||
self._request_timestamps, current_time
|
||||
)
|
||||
self._request_timestamps = timestamps
|
||||
|
||||
# Check if rate limit would be exceeded
|
||||
if len(timestamps) >= self.max_requests:
|
||||
# Need to wait
|
||||
pass
|
||||
else:
|
||||
# Slot available, record and proceed
|
||||
self._request_timestamps.append(current_time)
|
||||
return None
|
||||
|
||||
# Wait for availability if limit exceeded
|
||||
await self._wait_for_rate_limit(current_time)
|
||||
|
||||
# Record this request after waiting
|
||||
async with self._lock:
|
||||
current_time = time.time()
|
||||
self._request_timestamps.append(current_time)
|
||||
|
||||
# Allow request to proceed
|
||||
return None
|
||||
@@ -164,10 +164,6 @@ class AgentTool(BaseTool):
|
||||
if event.content:
|
||||
last_content = event.content
|
||||
|
||||
# Clean up runner resources (especially MCP sessions)
|
||||
# to avoid "Attempted to exit cancel scope in a different task" errors
|
||||
await runner.close()
|
||||
|
||||
if not last_content:
|
||||
return ''
|
||||
merged_text = '\n'.join(p.text for p in last_content.parts if p.text)
|
||||
|
||||
@@ -0,0 +1,239 @@
|
||||
# Copyright 2025 Google LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from unittest.mock import Mock
|
||||
|
||||
from google.adk.agents.callback_context import CallbackContext
|
||||
from google.adk.models.llm_request import LlmRequest
|
||||
from google.adk.plugins.rate_limit_plugin import RateLimitPlugin
|
||||
from google.genai import types
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def callback_context():
|
||||
"""Create a callback context for tests."""
|
||||
return Mock(spec=CallbackContext)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def llm_request():
|
||||
"""Create a basic LLM request for tests."""
|
||||
return LlmRequest(
|
||||
model='gemini-2.5-flash',
|
||||
contents=[
|
||||
types.Content(role='user', parts=[types.Part(text='Hello, world!')])
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
class TestRateLimitPlugin:
|
||||
"""Tests for the RateLimitPlugin."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_plugin_allows_requests_within_limit(
|
||||
self, callback_context, llm_request
|
||||
):
|
||||
"""Test that requests within the rate limit are allowed."""
|
||||
plugin = RateLimitPlugin(max_requests_per_minute=5)
|
||||
|
||||
# Send 5 requests (within limit)
|
||||
for _ in range(5):
|
||||
result = await plugin.before_model_callback(
|
||||
callback_context=callback_context, llm_request=llm_request
|
||||
)
|
||||
assert result is None # Request allowed
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_plugin_tracks_rate_limit_globally(self, callback_context):
|
||||
"""Test that rate limits are tracked globally across all models."""
|
||||
plugin = RateLimitPlugin(max_requests_per_minute=5)
|
||||
|
||||
# Create requests for two different models
|
||||
request_model1 = LlmRequest(
|
||||
model='gemini-2.5-flash',
|
||||
contents=[types.Content(role='user', parts=[types.Part(text='Hello')])],
|
||||
)
|
||||
request_model2 = LlmRequest(
|
||||
model='gemini-2.0-flash',
|
||||
contents=[types.Content(role='user', parts=[types.Part(text='Hello')])],
|
||||
)
|
||||
|
||||
# Send 3 requests to model1
|
||||
for _ in range(3):
|
||||
result = await plugin.before_model_callback(
|
||||
callback_context=callback_context, llm_request=request_model1
|
||||
)
|
||||
assert result is None
|
||||
|
||||
# Send 2 requests to model2 (fills global quota of 5)
|
||||
for _ in range(2):
|
||||
result = await plugin.before_model_callback(
|
||||
callback_context=callback_context, llm_request=request_model2
|
||||
)
|
||||
assert result is None
|
||||
|
||||
# Verify all 5 requests were tracked globally
|
||||
async with plugin._lock:
|
||||
assert len(plugin._request_timestamps) == 5
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_plugin_sliding_window_allows_requests_after_time(
|
||||
self, callback_context, llm_request
|
||||
):
|
||||
"""Test that requests are allowed after sliding window expires."""
|
||||
plugin = RateLimitPlugin(max_requests_per_minute=2)
|
||||
|
||||
# Send 2 requests (fill quota)
|
||||
for _ in range(2):
|
||||
result = await plugin.before_model_callback(
|
||||
callback_context=callback_context, llm_request=llm_request
|
||||
)
|
||||
assert result is None
|
||||
|
||||
# Simulate that requests happened 61 seconds ago (past the 60-second window)
|
||||
async with plugin._lock:
|
||||
plugin._request_timestamps = [
|
||||
time.time() - 61.0,
|
||||
time.time() - 61.0,
|
||||
]
|
||||
|
||||
# New request should be allowed immediately (old timestamps expired)
|
||||
result = await plugin.before_model_callback(
|
||||
callback_context=callback_context, llm_request=llm_request
|
||||
)
|
||||
assert result is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_plugin_waits_for_availability(
|
||||
self, callback_context, llm_request
|
||||
):
|
||||
"""Test that plugin waits when rate limit is exceeded."""
|
||||
plugin = RateLimitPlugin(max_requests_per_minute=2)
|
||||
|
||||
# Send 2 requests (fill quota)
|
||||
for _ in range(2):
|
||||
result = await plugin.before_model_callback(
|
||||
callback_context=callback_context, llm_request=llm_request
|
||||
)
|
||||
assert result is None
|
||||
|
||||
# Simulate old timestamps to avoid long wait
|
||||
async with plugin._lock:
|
||||
plugin._request_timestamps = [
|
||||
time.time() - 59.5, # Will expire in ~0.5 seconds
|
||||
time.time() - 59.5,
|
||||
]
|
||||
|
||||
# 3rd request should block briefly and then succeed
|
||||
start_time = time.time()
|
||||
result = await plugin.before_model_callback(
|
||||
callback_context=callback_context, llm_request=llm_request
|
||||
)
|
||||
elapsed_time = time.time() - start_time
|
||||
|
||||
assert result is None # Request eventually allowed
|
||||
assert elapsed_time >= 0.5 # Should have waited at least 0.5 seconds
|
||||
assert elapsed_time < 2.0 # But not too long
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_plugin_cleans_old_timestamps(
|
||||
self, callback_context, llm_request
|
||||
):
|
||||
"""Test that old timestamps are properly cleaned up."""
|
||||
plugin = RateLimitPlugin(max_requests_per_minute=5)
|
||||
|
||||
# Manually add old timestamps
|
||||
async with plugin._lock:
|
||||
plugin._request_timestamps = [
|
||||
time.time() - 120.0, # 2 minutes ago
|
||||
time.time() - 90.0, # 1.5 minutes ago
|
||||
time.time() - 70.0, # 70 seconds ago
|
||||
]
|
||||
|
||||
# New request should trigger cleanup and be allowed
|
||||
result = await plugin.before_model_callback(
|
||||
callback_context=callback_context, llm_request=llm_request
|
||||
)
|
||||
assert result is None
|
||||
|
||||
# Verify old timestamps were cleaned
|
||||
async with plugin._lock:
|
||||
timestamps = plugin._request_timestamps
|
||||
# Should only have the new timestamp
|
||||
assert len(timestamps) == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_plugin_concurrent_requests(self, callback_context):
|
||||
"""Test that plugin handles concurrent requests safely."""
|
||||
plugin = RateLimitPlugin(max_requests_per_minute=5)
|
||||
|
||||
# Create multiple requests (more than the limit)
|
||||
requests = [
|
||||
LlmRequest(
|
||||
model='gemini-2.5-flash',
|
||||
contents=[
|
||||
types.Content(
|
||||
role='user', parts=[types.Part(text=f'Request {i}')]
|
||||
)
|
||||
],
|
||||
)
|
||||
for i in range(7)
|
||||
]
|
||||
|
||||
# Simulate old timestamps for the first 5 requests to fill the quota
|
||||
async with plugin._lock:
|
||||
plugin._request_timestamps = [time.time() - 59.5 for _ in range(5)]
|
||||
|
||||
# Send 2 more requests concurrently (should both wait)
|
||||
tasks = [
|
||||
plugin.before_model_callback(
|
||||
callback_context=callback_context, llm_request=requests[i]
|
||||
)
|
||||
for i in range(2)
|
||||
]
|
||||
|
||||
start_time = time.time()
|
||||
results = await asyncio.gather(*tasks)
|
||||
elapsed_time = time.time() - start_time
|
||||
|
||||
# All requests should eventually succeed (after waiting)
|
||||
assert all(r is None for r in results)
|
||||
# Should have waited for at least one slot to become available
|
||||
assert elapsed_time >= 0.5
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_plugin_default_parameters(self, callback_context, llm_request):
|
||||
"""Test plugin works with default parameters."""
|
||||
plugin = RateLimitPlugin() # All defaults
|
||||
|
||||
# Should allow up to 15 requests (default max_requests_per_minute)
|
||||
for _ in range(15):
|
||||
result = await plugin.before_model_callback(
|
||||
callback_context=callback_context, llm_request=llm_request
|
||||
)
|
||||
assert result is None
|
||||
|
||||
# 16th request should block and wait
|
||||
# Simulate old timestamps to avoid long wait
|
||||
async with plugin._lock:
|
||||
plugin._request_timestamps = [time.time() - 59.8 for _ in range(15)]
|
||||
|
||||
# This should block briefly and succeed
|
||||
result = await plugin.before_model_callback(
|
||||
callback_context=callback_context, llm_request=llm_request
|
||||
)
|
||||
assert result is None
|
||||
@@ -131,10 +131,6 @@ async def test_agent_tool_inherits_parent_app_name(monkeypatch):
|
||||
)
|
||||
return _empty_async_generator()
|
||||
|
||||
async def close(self):
|
||||
"""Mock close method."""
|
||||
pass
|
||||
|
||||
monkeypatch.setattr('google.adk.runners.Runner', StubRunner)
|
||||
|
||||
tool_agent = Agent(
|
||||
|
||||
Reference in New Issue
Block a user