mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
fix: Support project-based gemini model path for BuiltInCodeExecutor and all built-in tools
This is to support model path like : projects/265104255505/locations/us-central1/publishers/google/models/gemini-2.0-flash-001" PiperOrigin-RevId: 783413351
This commit is contained in:
committed by
Copybara-Service
parent
30d7b37069
commit
a5d6f1e52e
@@ -12,11 +12,14 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from google.genai import types
|
||||
from typing_extensions import override
|
||||
|
||||
from ..agents.invocation_context import InvocationContext
|
||||
from ..models import LlmRequest
|
||||
from ..utils.model_name_utils import is_gemini_2_model
|
||||
from .base_code_executor import BaseCodeExecutor
|
||||
from .code_execution_utils import CodeExecutionInput
|
||||
from .code_execution_utils import CodeExecutionResult
|
||||
@@ -39,7 +42,7 @@ class BuiltInCodeExecutor(BaseCodeExecutor):
|
||||
|
||||
def process_llm_request(self, llm_request: LlmRequest) -> None:
|
||||
"""Pre-process the LLM request for Gemini 2.0+ models to use the code execution tool."""
|
||||
if llm_request.model and llm_request.model.startswith("gemini-2"):
|
||||
if is_gemini_2_model(llm_request.model):
|
||||
llm_request.config = llm_request.config or types.GenerateContentConfig()
|
||||
llm_request.config.tools = llm_request.config.tools or []
|
||||
llm_request.config.tools.append(
|
||||
|
||||
@@ -19,6 +19,8 @@ from typing import TYPE_CHECKING
|
||||
from google.genai import types
|
||||
from typing_extensions import override
|
||||
|
||||
from ..utils.model_name_utils import is_gemini_1_model
|
||||
from ..utils.model_name_utils import is_gemini_model
|
||||
from .base_tool import BaseTool
|
||||
from .tool_context import ToolContext
|
||||
|
||||
@@ -47,8 +49,8 @@ class EnterpriseWebSearchTool(BaseTool):
|
||||
tool_context: ToolContext,
|
||||
llm_request: LlmRequest,
|
||||
) -> None:
|
||||
if llm_request.model and 'gemini-' in llm_request.model:
|
||||
if 'gemini-1' in llm_request.model and llm_request.config.tools:
|
||||
if is_gemini_model(llm_request.model):
|
||||
if is_gemini_1_model(llm_request.model) and llm_request.config.tools:
|
||||
raise ValueError(
|
||||
'Enterprise web search tool can not be used with other tools in'
|
||||
' Gemini 1.x.'
|
||||
|
||||
@@ -19,6 +19,8 @@ from typing import TYPE_CHECKING
|
||||
from google.genai import types
|
||||
from typing_extensions import override
|
||||
|
||||
from ..utils.model_name_utils import is_gemini_1_model
|
||||
from ..utils.model_name_utils import is_gemini_model
|
||||
from .base_tool import BaseTool
|
||||
from .tool_context import ToolContext
|
||||
|
||||
@@ -46,7 +48,7 @@ class GoogleSearchTool(BaseTool):
|
||||
) -> None:
|
||||
llm_request.config = llm_request.config or types.GenerateContentConfig()
|
||||
llm_request.config.tools = llm_request.config.tools or []
|
||||
if llm_request.model and 'gemini-1' in llm_request.model:
|
||||
if is_gemini_1_model(llm_request.model):
|
||||
if llm_request.config.tools:
|
||||
raise ValueError(
|
||||
'Google search tool can not be used with other tools in Gemini 1.x.'
|
||||
@@ -54,7 +56,7 @@ class GoogleSearchTool(BaseTool):
|
||||
llm_request.config.tools.append(
|
||||
types.Tool(google_search_retrieval=types.GoogleSearchRetrieval())
|
||||
)
|
||||
elif llm_request.model and 'gemini-' in llm_request.model:
|
||||
elif is_gemini_model(llm_request.model):
|
||||
llm_request.config.tools.append(
|
||||
types.Tool(google_search=types.GoogleSearch())
|
||||
)
|
||||
|
||||
@@ -24,11 +24,12 @@ from google.genai import types
|
||||
from typing_extensions import override
|
||||
from vertexai.preview import rag
|
||||
|
||||
from ...utils.model_name_utils import is_gemini_2_model
|
||||
from ..tool_context import ToolContext
|
||||
from .base_retrieval_tool import BaseRetrievalTool
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ...models.llm_request import LlmRequest
|
||||
from ...models import LlmRequest
|
||||
|
||||
logger = logging.getLogger('google_adk.' + __name__)
|
||||
|
||||
@@ -62,7 +63,7 @@ class VertexAiRagRetrieval(BaseRetrievalTool):
|
||||
llm_request: LlmRequest,
|
||||
) -> None:
|
||||
# Use Gemini built-in Vertex AI RAG tool for Gemini 2 models.
|
||||
if llm_request.model and llm_request.model.startswith('gemini-2'):
|
||||
if is_gemini_2_model(llm_request.model):
|
||||
llm_request.config = (
|
||||
types.GenerateContentConfig()
|
||||
if not llm_request.config
|
||||
|
||||
@@ -19,6 +19,8 @@ from typing import TYPE_CHECKING
|
||||
from google.genai import types
|
||||
from typing_extensions import override
|
||||
|
||||
from ..utils.model_name_utils import is_gemini_1_model
|
||||
from ..utils.model_name_utils import is_gemini_2_model
|
||||
from .base_tool import BaseTool
|
||||
from .tool_context import ToolContext
|
||||
|
||||
@@ -46,9 +48,9 @@ class UrlContextTool(BaseTool):
|
||||
) -> None:
|
||||
llm_request.config = llm_request.config or types.GenerateContentConfig()
|
||||
llm_request.config.tools = llm_request.config.tools or []
|
||||
if llm_request.model and 'gemini-1' in llm_request.model:
|
||||
if is_gemini_1_model(llm_request.model):
|
||||
raise ValueError('Url context tool can not be used in Gemini 1.x.')
|
||||
elif llm_request.model and 'gemini-2' in llm_request.model:
|
||||
elif is_gemini_2_model(llm_request.model):
|
||||
llm_request.config.tools.append(
|
||||
types.Tool(url_context=types.UrlContext())
|
||||
)
|
||||
|
||||
@@ -20,6 +20,8 @@ from typing import TYPE_CHECKING
|
||||
from google.genai import types
|
||||
from typing_extensions import override
|
||||
|
||||
from ..utils.model_name_utils import is_gemini_1_model
|
||||
from ..utils.model_name_utils import is_gemini_model
|
||||
from .base_tool import BaseTool
|
||||
from .tool_context import ToolContext
|
||||
|
||||
@@ -86,8 +88,8 @@ class VertexAiSearchTool(BaseTool):
|
||||
tool_context: ToolContext,
|
||||
llm_request: LlmRequest,
|
||||
) -> None:
|
||||
if llm_request.model and 'gemini-' in llm_request.model:
|
||||
if 'gemini-1' in llm_request.model and llm_request.config.tools:
|
||||
if is_gemini_model(llm_request.model):
|
||||
if is_gemini_1_model(llm_request.model) and llm_request.config.tools:
|
||||
raise ValueError(
|
||||
'Vertex AI search tool can not be used with other tools in Gemini'
|
||||
' 1.x.'
|
||||
|
||||
@@ -0,0 +1,90 @@
|
||||
# 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.
|
||||
|
||||
"""Utilities for model name validation and parsing."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from typing import Optional
|
||||
|
||||
|
||||
def extract_model_name(model_string: str) -> str:
|
||||
"""Extract the actual model name from either simple or path-based format.
|
||||
|
||||
Args:
|
||||
model_string: Either a simple model name like "gemini-2.5-pro" or
|
||||
a path-based model name like "projects/.../models/gemini-2.0-flash-001"
|
||||
|
||||
Returns:
|
||||
The extracted model name (e.g., "gemini-2.5-pro")
|
||||
"""
|
||||
# Pattern for path-based model names
|
||||
path_pattern = (
|
||||
r'^projects/[^/]+/locations/[^/]+/publishers/[^/]+/models/(.+)$'
|
||||
)
|
||||
match = re.match(path_pattern, model_string)
|
||||
if match:
|
||||
return match.group(1)
|
||||
|
||||
# If it's not a path-based model, return as-is (simple model name)
|
||||
return model_string
|
||||
|
||||
|
||||
def is_gemini_model(model_string: Optional[str]) -> bool:
|
||||
"""Check if the model is a Gemini model using regex patterns.
|
||||
|
||||
Args:
|
||||
model_string: Either a simple model name or path-based model name
|
||||
|
||||
Returns:
|
||||
True if it's a Gemini model, False otherwise
|
||||
"""
|
||||
if not model_string:
|
||||
return False
|
||||
|
||||
model_name = extract_model_name(model_string)
|
||||
return re.match(r'^gemini-', model_name) is not None
|
||||
|
||||
|
||||
def is_gemini_1_model(model_string: Optional[str]) -> bool:
|
||||
"""Check if the model is a Gemini 1.x model using regex patterns.
|
||||
|
||||
Args:
|
||||
model_string: Either a simple model name or path-based model name
|
||||
|
||||
Returns:
|
||||
True if it's a Gemini 1.x model, False otherwise
|
||||
"""
|
||||
if not model_string:
|
||||
return False
|
||||
|
||||
model_name = extract_model_name(model_string)
|
||||
return re.match(r'^gemini-1\.\d+', model_name) is not None
|
||||
|
||||
|
||||
def is_gemini_2_model(model_string: Optional[str]) -> bool:
|
||||
"""Check if the model is a Gemini 2.x model using regex patterns.
|
||||
|
||||
Args:
|
||||
model_string: Either a simple model name or path-based model name
|
||||
|
||||
Returns:
|
||||
True if it's a Gemini 2.x model, False otherwise
|
||||
"""
|
||||
if not model_string:
|
||||
return False
|
||||
|
||||
model_name = extract_model_name(model_string)
|
||||
return re.match(r'^gemini-2\.\d+', model_name) is not None
|
||||
@@ -0,0 +1,434 @@
|
||||
# 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.
|
||||
|
||||
"""Tests for GoogleSearchTool."""
|
||||
|
||||
from google.adk.agents.invocation_context import InvocationContext
|
||||
from google.adk.agents.sequential_agent import SequentialAgent
|
||||
from google.adk.models.llm_request import LlmRequest
|
||||
from google.adk.sessions.in_memory_session_service import InMemorySessionService
|
||||
from google.adk.tools.google_search_tool import google_search
|
||||
from google.adk.tools.google_search_tool import GoogleSearchTool
|
||||
from google.adk.tools.tool_context import ToolContext
|
||||
from google.genai import types
|
||||
import pytest
|
||||
|
||||
|
||||
async def _create_tool_context() -> ToolContext:
|
||||
session_service = InMemorySessionService()
|
||||
session = await session_service.create_session(
|
||||
app_name='test_app', user_id='test_user'
|
||||
)
|
||||
agent = SequentialAgent(name='test_agent')
|
||||
invocation_context = InvocationContext(
|
||||
invocation_id='invocation_id',
|
||||
agent=agent,
|
||||
session=session,
|
||||
session_service=session_service,
|
||||
)
|
||||
return ToolContext(invocation_context=invocation_context)
|
||||
|
||||
|
||||
class TestGoogleSearchTool:
|
||||
"""Test the GoogleSearchTool class."""
|
||||
|
||||
def test_init(self):
|
||||
"""Test initialization of GoogleSearchTool."""
|
||||
tool = GoogleSearchTool()
|
||||
assert tool.name == 'google_search'
|
||||
assert tool.description == 'google_search'
|
||||
|
||||
def test_google_search_singleton(self):
|
||||
"""Test that google_search is a singleton instance."""
|
||||
assert isinstance(google_search, GoogleSearchTool)
|
||||
assert google_search.name == 'google_search'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_llm_request_with_gemini_1_model(self):
|
||||
"""Test processing LLM request with Gemini 1.x model."""
|
||||
tool = GoogleSearchTool()
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
llm_request = LlmRequest(
|
||||
model='gemini-1.5-flash', config=types.GenerateContentConfig()
|
||||
)
|
||||
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
assert llm_request.config.tools is not None
|
||||
assert len(llm_request.config.tools) == 1
|
||||
assert llm_request.config.tools[0].google_search_retrieval is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_llm_request_with_path_based_gemini_1_model(self):
|
||||
"""Test processing LLM request with path-based Gemini 1.x model."""
|
||||
tool = GoogleSearchTool()
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
llm_request = LlmRequest(
|
||||
model='projects/265104255505/locations/us-central1/publishers/google/models/gemini-1.5-flash-001',
|
||||
config=types.GenerateContentConfig(),
|
||||
)
|
||||
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
assert llm_request.config.tools is not None
|
||||
assert len(llm_request.config.tools) == 1
|
||||
assert llm_request.config.tools[0].google_search_retrieval is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_llm_request_with_gemini_1_0_model(self):
|
||||
"""Test processing LLM request with Gemini 1.0 model."""
|
||||
tool = GoogleSearchTool()
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
llm_request = LlmRequest(
|
||||
model='gemini-1.0-pro', config=types.GenerateContentConfig()
|
||||
)
|
||||
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
assert llm_request.config.tools is not None
|
||||
assert len(llm_request.config.tools) == 1
|
||||
assert llm_request.config.tools[0].google_search_retrieval is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_llm_request_with_gemini_2_model(self):
|
||||
"""Test processing LLM request with Gemini 2.x model."""
|
||||
tool = GoogleSearchTool()
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
llm_request = LlmRequest(
|
||||
model='gemini-2.0-flash', config=types.GenerateContentConfig()
|
||||
)
|
||||
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
assert llm_request.config.tools is not None
|
||||
assert len(llm_request.config.tools) == 1
|
||||
assert llm_request.config.tools[0].google_search is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_llm_request_with_path_based_gemini_2_model(self):
|
||||
"""Test processing LLM request with path-based Gemini 2.x model."""
|
||||
tool = GoogleSearchTool()
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
llm_request = LlmRequest(
|
||||
model='projects/265104255505/locations/us-central1/publishers/google/models/gemini-2.0-flash-001',
|
||||
config=types.GenerateContentConfig(),
|
||||
)
|
||||
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
assert llm_request.config.tools is not None
|
||||
assert len(llm_request.config.tools) == 1
|
||||
assert llm_request.config.tools[0].google_search is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_llm_request_with_gemini_2_5_model(self):
|
||||
"""Test processing LLM request with Gemini 2.5 model."""
|
||||
tool = GoogleSearchTool()
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
llm_request = LlmRequest(
|
||||
model='gemini-2.5-pro', config=types.GenerateContentConfig()
|
||||
)
|
||||
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
assert llm_request.config.tools is not None
|
||||
assert len(llm_request.config.tools) == 1
|
||||
assert llm_request.config.tools[0].google_search is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_llm_request_with_gemini_1_model_and_existing_tools_raises_error(
|
||||
self,
|
||||
):
|
||||
"""Test that Gemini 1.x model with existing tools raises ValueError."""
|
||||
tool = GoogleSearchTool()
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
existing_tool = types.Tool(
|
||||
function_declarations=[
|
||||
types.FunctionDeclaration(name='test_function', description='test')
|
||||
]
|
||||
)
|
||||
|
||||
llm_request = LlmRequest(
|
||||
model='gemini-1.5-flash',
|
||||
config=types.GenerateContentConfig(tools=[existing_tool]),
|
||||
)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=(
|
||||
'Google search tool can not be used with other tools in Gemini 1.x'
|
||||
),
|
||||
):
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_llm_request_with_path_based_gemini_1_model_and_existing_tools_raises_error(
|
||||
self,
|
||||
):
|
||||
"""Test that path-based Gemini 1.x model with existing tools raises ValueError."""
|
||||
tool = GoogleSearchTool()
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
existing_tool = types.Tool(
|
||||
function_declarations=[
|
||||
types.FunctionDeclaration(name='test_function', description='test')
|
||||
]
|
||||
)
|
||||
|
||||
llm_request = LlmRequest(
|
||||
model='projects/265104255505/locations/us-central1/publishers/google/models/gemini-1.5-pro-preview',
|
||||
config=types.GenerateContentConfig(tools=[existing_tool]),
|
||||
)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=(
|
||||
'Google search tool can not be used with other tools in Gemini 1.x'
|
||||
),
|
||||
):
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_llm_request_with_gemini_2_model_and_existing_tools_succeeds(
|
||||
self,
|
||||
):
|
||||
"""Test that Gemini 2.x model with existing tools succeeds."""
|
||||
tool = GoogleSearchTool()
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
existing_tool = types.Tool(
|
||||
function_declarations=[
|
||||
types.FunctionDeclaration(name='test_function', description='test')
|
||||
]
|
||||
)
|
||||
|
||||
llm_request = LlmRequest(
|
||||
model='gemini-2.0-flash',
|
||||
config=types.GenerateContentConfig(tools=[existing_tool]),
|
||||
)
|
||||
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
assert llm_request.config.tools is not None
|
||||
assert len(llm_request.config.tools) == 2
|
||||
assert llm_request.config.tools[0] == existing_tool
|
||||
assert llm_request.config.tools[1].google_search is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_llm_request_with_non_gemini_model_raises_error(self):
|
||||
"""Test that non-Gemini model raises ValueError."""
|
||||
tool = GoogleSearchTool()
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
llm_request = LlmRequest(
|
||||
model='claude-3-sonnet', config=types.GenerateContentConfig()
|
||||
)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match='Google search tool is not supported for model claude-3-sonnet',
|
||||
):
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_llm_request_with_path_based_non_gemini_model_raises_error(
|
||||
self,
|
||||
):
|
||||
"""Test that path-based non-Gemini model raises ValueError."""
|
||||
tool = GoogleSearchTool()
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
non_gemini_path = 'projects/265104255505/locations/us-central1/publishers/google/models/claude-3-sonnet'
|
||||
llm_request = LlmRequest(
|
||||
model=non_gemini_path, config=types.GenerateContentConfig()
|
||||
)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=(
|
||||
f'Google search tool is not supported for model {non_gemini_path}'
|
||||
),
|
||||
):
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_llm_request_with_none_model_raises_error(self):
|
||||
"""Test that None model raises ValueError."""
|
||||
tool = GoogleSearchTool()
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
llm_request = LlmRequest(model=None, config=types.GenerateContentConfig())
|
||||
|
||||
with pytest.raises(
|
||||
ValueError, match='Google search tool is not supported for model None'
|
||||
):
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_llm_request_with_empty_model_raises_error(self):
|
||||
"""Test that empty model raises ValueError."""
|
||||
tool = GoogleSearchTool()
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
llm_request = LlmRequest(model='', config=types.GenerateContentConfig())
|
||||
|
||||
with pytest.raises(
|
||||
ValueError, match='Google search tool is not supported for model '
|
||||
):
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_llm_request_with_none_config(self):
|
||||
"""Test processing LLM request with None config."""
|
||||
tool = GoogleSearchTool()
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
llm_request = LlmRequest(model='gemini-2.0-flash', config=None)
|
||||
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
assert llm_request.config is not None
|
||||
assert llm_request.config.tools is not None
|
||||
assert len(llm_request.config.tools) == 1
|
||||
assert llm_request.config.tools[0].google_search is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_llm_request_with_none_tools(self):
|
||||
"""Test processing LLM request with None tools."""
|
||||
tool = GoogleSearchTool()
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
llm_request = LlmRequest(
|
||||
model='gemini-2.0-flash', config=types.GenerateContentConfig(tools=None)
|
||||
)
|
||||
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
assert llm_request.config.tools is not None
|
||||
assert len(llm_request.config.tools) == 1
|
||||
assert llm_request.config.tools[0].google_search is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_llm_request_edge_cases(self):
|
||||
"""Test edge cases for model name validation."""
|
||||
tool = GoogleSearchTool()
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
# Test with model names that contain gemini but don't start with it
|
||||
edge_cases = [
|
||||
'my-gemini-1.5-model',
|
||||
'custom-gemini-2.0-flash',
|
||||
'projects/265104255505/locations/us-central1/publishers/gemini/models/claude-3-sonnet',
|
||||
]
|
||||
|
||||
for model in edge_cases:
|
||||
llm_request = LlmRequest(
|
||||
model=model, config=types.GenerateContentConfig()
|
||||
)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=f'Google search tool is not supported for model {model}',
|
||||
):
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_llm_request_gemini_version_specifics(self):
|
||||
"""Test specific Gemini version behaviors."""
|
||||
tool = GoogleSearchTool()
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
# Test various Gemini versions
|
||||
gemini_1_models = [
|
||||
'gemini-1.0-pro',
|
||||
'gemini-1.5-flash',
|
||||
'gemini-1.5-pro',
|
||||
'gemini-1.9-experimental',
|
||||
]
|
||||
|
||||
gemini_2_models = [
|
||||
'gemini-2.0-flash',
|
||||
'gemini-2.0-pro',
|
||||
'gemini-2.5-flash',
|
||||
'gemini-2.5-pro',
|
||||
]
|
||||
|
||||
# Test Gemini 1.x models use google_search_retrieval
|
||||
for model in gemini_1_models:
|
||||
llm_request = LlmRequest(
|
||||
model=model, config=types.GenerateContentConfig()
|
||||
)
|
||||
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
assert llm_request.config.tools is not None
|
||||
assert len(llm_request.config.tools) == 1
|
||||
assert llm_request.config.tools[0].google_search_retrieval is not None
|
||||
assert llm_request.config.tools[0].google_search is None
|
||||
|
||||
# Test Gemini 2.x models use google_search
|
||||
for model in gemini_2_models:
|
||||
llm_request = LlmRequest(
|
||||
model=model, config=types.GenerateContentConfig()
|
||||
)
|
||||
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
assert llm_request.config.tools is not None
|
||||
assert len(llm_request.config.tools) == 1
|
||||
assert llm_request.config.tools[0].google_search is not None
|
||||
assert llm_request.config.tools[0].google_search_retrieval is None
|
||||
@@ -0,0 +1,303 @@
|
||||
# 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.
|
||||
|
||||
"""Tests for UrlContextTool."""
|
||||
|
||||
from google.adk.agents.invocation_context import InvocationContext
|
||||
from google.adk.agents.sequential_agent import SequentialAgent
|
||||
from google.adk.models.llm_request import LlmRequest
|
||||
from google.adk.sessions.in_memory_session_service import InMemorySessionService
|
||||
from google.adk.tools.tool_context import ToolContext
|
||||
from google.adk.tools.url_context_tool import url_context
|
||||
from google.adk.tools.url_context_tool import UrlContextTool
|
||||
from google.genai import types
|
||||
import pytest
|
||||
|
||||
|
||||
async def _create_tool_context() -> ToolContext:
|
||||
session_service = InMemorySessionService()
|
||||
session = await session_service.create_session(
|
||||
app_name='test_app', user_id='test_user'
|
||||
)
|
||||
agent = SequentialAgent(name='test_agent')
|
||||
invocation_context = InvocationContext(
|
||||
invocation_id='invocation_id',
|
||||
agent=agent,
|
||||
session=session,
|
||||
session_service=session_service,
|
||||
)
|
||||
return ToolContext(invocation_context=invocation_context)
|
||||
|
||||
|
||||
class TestUrlContextTool:
|
||||
"""Test the UrlContextTool class."""
|
||||
|
||||
def test_init(self):
|
||||
"""Test initialization of UrlContextTool."""
|
||||
tool = UrlContextTool()
|
||||
assert tool.name == 'url_context'
|
||||
assert tool.description == 'url_context'
|
||||
|
||||
def test_url_context_singleton(self):
|
||||
"""Test that url_context is a singleton instance."""
|
||||
assert isinstance(url_context, UrlContextTool)
|
||||
assert url_context.name == 'url_context'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_llm_request_with_gemini_2_model(self):
|
||||
"""Test processing LLM request with Gemini 2.x model."""
|
||||
tool = UrlContextTool()
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
llm_request = LlmRequest(
|
||||
model='gemini-2.0-flash', config=types.GenerateContentConfig()
|
||||
)
|
||||
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
assert llm_request.config.tools is not None
|
||||
assert len(llm_request.config.tools) == 1
|
||||
assert llm_request.config.tools[0].url_context is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_llm_request_with_path_based_gemini_2_model(self):
|
||||
"""Test processing LLM request with path-based Gemini 2.x model."""
|
||||
tool = UrlContextTool()
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
llm_request = LlmRequest(
|
||||
model='projects/265104255505/locations/us-central1/publishers/google/models/gemini-2.0-flash-001',
|
||||
config=types.GenerateContentConfig(),
|
||||
)
|
||||
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
assert llm_request.config.tools is not None
|
||||
assert len(llm_request.config.tools) == 1
|
||||
assert llm_request.config.tools[0].url_context is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_llm_request_with_gemini_2_5_model(self):
|
||||
"""Test processing LLM request with Gemini 2.5 model."""
|
||||
tool = UrlContextTool()
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
llm_request = LlmRequest(
|
||||
model='gemini-2.5-pro', config=types.GenerateContentConfig()
|
||||
)
|
||||
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
assert llm_request.config.tools is not None
|
||||
assert len(llm_request.config.tools) == 1
|
||||
assert llm_request.config.tools[0].url_context is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_llm_request_with_existing_tools(self):
|
||||
"""Test processing LLM request with existing tools."""
|
||||
tool = UrlContextTool()
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
existing_tool = types.Tool(
|
||||
function_declarations=[
|
||||
types.FunctionDeclaration(name='test_function', description='test')
|
||||
]
|
||||
)
|
||||
|
||||
llm_request = LlmRequest(
|
||||
model='gemini-2.0-flash',
|
||||
config=types.GenerateContentConfig(tools=[existing_tool]),
|
||||
)
|
||||
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
assert llm_request.config.tools is not None
|
||||
assert len(llm_request.config.tools) == 2
|
||||
assert llm_request.config.tools[0] == existing_tool
|
||||
assert llm_request.config.tools[1].url_context is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_llm_request_with_gemini_1_model_raises_error(self):
|
||||
"""Test that Gemini 1.x model raises ValueError."""
|
||||
tool = UrlContextTool()
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
llm_request = LlmRequest(
|
||||
model='gemini-1.5-flash', config=types.GenerateContentConfig()
|
||||
)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError, match='Url context tool can not be used in Gemini 1.x'
|
||||
):
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_llm_request_with_path_based_gemini_1_model_raises_error(
|
||||
self,
|
||||
):
|
||||
"""Test that path-based Gemini 1.x model raises ValueError."""
|
||||
tool = UrlContextTool()
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
llm_request = LlmRequest(
|
||||
model='projects/265104255505/locations/us-central1/publishers/google/models/gemini-1.5-flash-001',
|
||||
config=types.GenerateContentConfig(),
|
||||
)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError, match='Url context tool can not be used in Gemini 1.x'
|
||||
):
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_llm_request_with_non_gemini_model_raises_error(self):
|
||||
"""Test that non-Gemini model raises ValueError."""
|
||||
tool = UrlContextTool()
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
llm_request = LlmRequest(
|
||||
model='claude-3-sonnet', config=types.GenerateContentConfig()
|
||||
)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match='Url context tool is not supported for model claude-3-sonnet',
|
||||
):
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_llm_request_with_path_based_non_gemini_model_raises_error(
|
||||
self,
|
||||
):
|
||||
"""Test that path-based non-Gemini model raises ValueError."""
|
||||
tool = UrlContextTool()
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
non_gemini_path = 'projects/265104255505/locations/us-central1/publishers/google/models/claude-3-sonnet'
|
||||
llm_request = LlmRequest(
|
||||
model=non_gemini_path, config=types.GenerateContentConfig()
|
||||
)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=f'Url context tool is not supported for model {non_gemini_path}',
|
||||
):
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_llm_request_with_none_model_raises_error(self):
|
||||
"""Test that None model raises ValueError."""
|
||||
tool = UrlContextTool()
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
llm_request = LlmRequest(model=None, config=types.GenerateContentConfig())
|
||||
|
||||
with pytest.raises(
|
||||
ValueError, match='Url context tool is not supported for model None'
|
||||
):
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_llm_request_with_empty_model_raises_error(self):
|
||||
"""Test that empty model raises ValueError."""
|
||||
tool = UrlContextTool()
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
llm_request = LlmRequest(model='', config=types.GenerateContentConfig())
|
||||
|
||||
with pytest.raises(
|
||||
ValueError, match='Url context tool is not supported for model '
|
||||
):
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_llm_request_with_none_config(self):
|
||||
"""Test processing LLM request with None config."""
|
||||
tool = UrlContextTool()
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
llm_request = LlmRequest(model='gemini-2.0-flash', config=None)
|
||||
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
assert llm_request.config is not None
|
||||
assert llm_request.config.tools is not None
|
||||
assert len(llm_request.config.tools) == 1
|
||||
assert llm_request.config.tools[0].url_context is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_llm_request_with_none_tools(self):
|
||||
"""Test processing LLM request with None tools."""
|
||||
tool = UrlContextTool()
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
llm_request = LlmRequest(
|
||||
model='gemini-2.0-flash', config=types.GenerateContentConfig(tools=None)
|
||||
)
|
||||
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
assert llm_request.config.tools is not None
|
||||
assert len(llm_request.config.tools) == 1
|
||||
assert llm_request.config.tools[0].url_context is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_llm_request_edge_cases(self):
|
||||
"""Test edge cases for model name validation."""
|
||||
tool = UrlContextTool()
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
# Test with model names that contain gemini but don't start with it
|
||||
edge_cases = [
|
||||
'my-gemini-2.0-model',
|
||||
'custom-gemini-2.5-flash',
|
||||
'projects/265104255505/locations/us-central1/publishers/gemini/models/claude-3-sonnet',
|
||||
]
|
||||
|
||||
for model in edge_cases:
|
||||
llm_request = LlmRequest(
|
||||
model=model, config=types.GenerateContentConfig()
|
||||
)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=f'Url context tool is not supported for model {model}',
|
||||
):
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
@@ -0,0 +1,320 @@
|
||||
# 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.invocation_context import InvocationContext
|
||||
from google.adk.agents.sequential_agent import SequentialAgent
|
||||
from google.adk.models.llm_request import LlmRequest
|
||||
from google.adk.sessions.in_memory_session_service import InMemorySessionService
|
||||
from google.adk.tools.tool_context import ToolContext
|
||||
from google.adk.tools.vertex_ai_search_tool import VertexAiSearchTool
|
||||
from google.adk.utils.model_name_utils import extract_model_name
|
||||
from google.adk.utils.model_name_utils import is_gemini_1_model
|
||||
from google.adk.utils.model_name_utils import is_gemini_model
|
||||
from google.genai import types
|
||||
import pytest
|
||||
|
||||
|
||||
async def _create_tool_context() -> ToolContext:
|
||||
session_service = InMemorySessionService()
|
||||
session = await session_service.create_session(
|
||||
app_name='test_app', user_id='test_user'
|
||||
)
|
||||
agent = SequentialAgent(name='test_agent')
|
||||
invocation_context = InvocationContext(
|
||||
invocation_id='invocation_id',
|
||||
agent=agent,
|
||||
session=session,
|
||||
session_service=session_service,
|
||||
)
|
||||
return ToolContext(invocation_context=invocation_context)
|
||||
|
||||
|
||||
class TestVertexAiSearchToolHelperFunctions:
|
||||
"""Test the helper functions for model name extraction and validation."""
|
||||
|
||||
def test_extract_model_name_simple_model(self):
|
||||
"""Test extraction of simple model names."""
|
||||
assert extract_model_name('gemini-2.5-pro') == 'gemini-2.5-pro'
|
||||
assert extract_model_name('gemini-1.5-flash') == 'gemini-1.5-flash'
|
||||
assert extract_model_name('gemini-1.0-pro') == 'gemini-1.0-pro'
|
||||
assert extract_model_name('claude-3-sonnet') == 'claude-3-sonnet'
|
||||
|
||||
def test_extract_model_name_path_based_model(self):
|
||||
"""Test extraction of path-based model names."""
|
||||
path_model = 'projects/265104255505/locations/us-central1/publishers/google/models/gemini-2.0-flash-001'
|
||||
assert extract_model_name(path_model) == 'gemini-2.0-flash-001'
|
||||
|
||||
path_model_2 = 'projects/12345/locations/us-east1/publishers/google/models/gemini-1.5-pro-preview'
|
||||
assert extract_model_name(path_model_2) == 'gemini-1.5-pro-preview'
|
||||
|
||||
def test_extract_model_name_invalid_path(self):
|
||||
"""Test that invalid path formats return the original string."""
|
||||
invalid_path = 'projects/invalid/path/format'
|
||||
assert extract_model_name(invalid_path) == invalid_path
|
||||
|
||||
def test_is_gemini_model_simple_names(self):
|
||||
"""Test Gemini model detection with simple model names."""
|
||||
assert is_gemini_model('gemini-2.5-pro') is True
|
||||
assert is_gemini_model('gemini-1.5-flash') is True
|
||||
assert is_gemini_model('gemini-1.0-pro') is True
|
||||
assert is_gemini_model('claude-3-sonnet') is False
|
||||
assert is_gemini_model('gpt-4') is False
|
||||
assert is_gemini_model('gemini') is False # Must have dash after gemini
|
||||
|
||||
def test_is_gemini_model_path_based_names(self):
|
||||
"""Test Gemini model detection with path-based model names."""
|
||||
gemini_path = 'projects/265104255505/locations/us-central1/publishers/google/models/gemini-2.0-flash-001'
|
||||
assert is_gemini_model(gemini_path) is True
|
||||
|
||||
non_gemini_path = 'projects/265104255505/locations/us-central1/publishers/google/models/claude-3-sonnet'
|
||||
assert is_gemini_model(non_gemini_path) is False
|
||||
|
||||
def test_is_gemini_1_model_simple_names(self):
|
||||
"""Test Gemini 1.x model detection with simple model names."""
|
||||
assert is_gemini_1_model('gemini-1.5-flash') is True
|
||||
assert is_gemini_1_model('gemini-1.0-pro') is True
|
||||
assert is_gemini_1_model('gemini-1.5-pro-preview') is True
|
||||
assert is_gemini_1_model('gemini-2.0-flash') is False
|
||||
assert is_gemini_1_model('gemini-2.5-pro') is False
|
||||
assert is_gemini_1_model('gemini-10.0-pro') is False # Only 1.x versions
|
||||
assert is_gemini_1_model('claude-3-sonnet') is False
|
||||
|
||||
def test_is_gemini_1_model_path_based_names(self):
|
||||
"""Test Gemini 1.x model detection with path-based model names."""
|
||||
gemini_1_path = 'projects/265104255505/locations/us-central1/publishers/google/models/gemini-1.5-flash-001'
|
||||
assert is_gemini_1_model(gemini_1_path) is True
|
||||
|
||||
gemini_2_path = 'projects/265104255505/locations/us-central1/publishers/google/models/gemini-2.0-flash-001'
|
||||
assert is_gemini_1_model(gemini_2_path) is False
|
||||
|
||||
def test_edge_cases(self):
|
||||
"""Test edge cases for model name validation."""
|
||||
# Test with empty string
|
||||
assert is_gemini_model('') is False
|
||||
assert is_gemini_1_model('') is False
|
||||
|
||||
# Test with model names containing gemini but not starting with it
|
||||
assert is_gemini_model('my-gemini-model') is False
|
||||
assert is_gemini_1_model('my-gemini-1.5-model') is False
|
||||
|
||||
# Test with model names that have gemini in the middle of the path
|
||||
tricky_path = 'projects/265104255505/locations/us-central1/publishers/gemini/models/claude-3-sonnet'
|
||||
assert is_gemini_model(tricky_path) is False
|
||||
|
||||
|
||||
class TestVertexAiSearchTool:
|
||||
"""Test the VertexAiSearchTool class."""
|
||||
|
||||
def test_init_with_data_store_id(self):
|
||||
"""Test initialization with data store ID."""
|
||||
tool = VertexAiSearchTool(data_store_id='test_data_store')
|
||||
assert tool.data_store_id == 'test_data_store'
|
||||
assert tool.search_engine_id is None
|
||||
|
||||
def test_init_with_search_engine_id(self):
|
||||
"""Test initialization with search engine ID."""
|
||||
tool = VertexAiSearchTool(search_engine_id='test_search_engine')
|
||||
assert tool.search_engine_id == 'test_search_engine'
|
||||
assert tool.data_store_id is None
|
||||
|
||||
def test_init_with_neither_raises_error(self):
|
||||
"""Test that initialization without either ID raises ValueError."""
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match='Either data_store_id or search_engine_id must be specified',
|
||||
):
|
||||
VertexAiSearchTool()
|
||||
|
||||
def test_init_with_both_raises_error(self):
|
||||
"""Test that initialization with both IDs raises ValueError."""
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match='Either data_store_id or search_engine_id must be specified',
|
||||
):
|
||||
VertexAiSearchTool(
|
||||
data_store_id='test_data_store', search_engine_id='test_search_engine'
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_llm_request_with_simple_gemini_model(self):
|
||||
"""Test processing LLM request with simple Gemini model name."""
|
||||
tool = VertexAiSearchTool(data_store_id='test_data_store')
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
llm_request = LlmRequest(
|
||||
model='gemini-2.5-pro', config=types.GenerateContentConfig()
|
||||
)
|
||||
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
assert llm_request.config.tools is not None
|
||||
assert len(llm_request.config.tools) == 1
|
||||
assert llm_request.config.tools[0].retrieval is not None
|
||||
assert llm_request.config.tools[0].retrieval.vertex_ai_search is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_llm_request_with_path_based_gemini_model(self):
|
||||
"""Test processing LLM request with path-based Gemini model name."""
|
||||
tool = VertexAiSearchTool(data_store_id='test_data_store')
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
llm_request = LlmRequest(
|
||||
model='projects/265104255505/locations/us-central1/publishers/google/models/gemini-2.0-flash-001',
|
||||
config=types.GenerateContentConfig(),
|
||||
)
|
||||
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
assert llm_request.config.tools is not None
|
||||
assert len(llm_request.config.tools) == 1
|
||||
assert llm_request.config.tools[0].retrieval is not None
|
||||
assert llm_request.config.tools[0].retrieval.vertex_ai_search is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_llm_request_with_gemini_1_and_other_tools_raises_error(
|
||||
self,
|
||||
):
|
||||
"""Test that Gemini 1.x with other tools raises ValueError."""
|
||||
tool = VertexAiSearchTool(data_store_id='test_data_store')
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
existing_tool = types.Tool(
|
||||
function_declarations=[
|
||||
types.FunctionDeclaration(name='test_function', description='test')
|
||||
]
|
||||
)
|
||||
|
||||
llm_request = LlmRequest(
|
||||
model='gemini-1.5-flash',
|
||||
config=types.GenerateContentConfig(tools=[existing_tool]),
|
||||
)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=(
|
||||
'Vertex AI search tool can not be used with other tools in'
|
||||
' Gemini 1.x'
|
||||
),
|
||||
):
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_llm_request_with_path_based_gemini_1_and_other_tools_raises_error(
|
||||
self,
|
||||
):
|
||||
"""Test that path-based Gemini 1.x with other tools raises ValueError."""
|
||||
tool = VertexAiSearchTool(data_store_id='test_data_store')
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
existing_tool = types.Tool(
|
||||
function_declarations=[
|
||||
types.FunctionDeclaration(name='test_function', description='test')
|
||||
]
|
||||
)
|
||||
|
||||
llm_request = LlmRequest(
|
||||
model='projects/265104255505/locations/us-central1/publishers/google/models/gemini-1.5-pro-preview',
|
||||
config=types.GenerateContentConfig(tools=[existing_tool]),
|
||||
)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=(
|
||||
'Vertex AI search tool can not be used with other tools in'
|
||||
' Gemini 1.x'
|
||||
),
|
||||
):
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_llm_request_with_non_gemini_model_raises_error(self):
|
||||
"""Test that non-Gemini model raises ValueError."""
|
||||
tool = VertexAiSearchTool(data_store_id='test_data_store')
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
llm_request = LlmRequest(
|
||||
model='claude-3-sonnet', config=types.GenerateContentConfig()
|
||||
)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=(
|
||||
'Vertex AI search tool is not supported for model claude-3-sonnet'
|
||||
),
|
||||
):
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_llm_request_with_path_based_non_gemini_model_raises_error(
|
||||
self,
|
||||
):
|
||||
"""Test that path-based non-Gemini model raises ValueError."""
|
||||
tool = VertexAiSearchTool(data_store_id='test_data_store')
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
non_gemini_path = 'projects/265104255505/locations/us-central1/publishers/google/models/claude-3-sonnet'
|
||||
llm_request = LlmRequest(
|
||||
model=non_gemini_path, config=types.GenerateContentConfig()
|
||||
)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=(
|
||||
'Vertex AI search tool is not supported for model'
|
||||
f' {non_gemini_path}'
|
||||
),
|
||||
):
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_llm_request_with_gemini_2_and_other_tools_succeeds(
|
||||
self,
|
||||
):
|
||||
"""Test that Gemini 2.x with other tools succeeds."""
|
||||
tool = VertexAiSearchTool(data_store_id='test_data_store')
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
existing_tool = types.Tool(
|
||||
function_declarations=[
|
||||
types.FunctionDeclaration(name='test_function', description='test')
|
||||
]
|
||||
)
|
||||
|
||||
llm_request = LlmRequest(
|
||||
model='gemini-2.5-pro',
|
||||
config=types.GenerateContentConfig(tools=[existing_tool]),
|
||||
)
|
||||
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
# Should have both the existing tool and the new vertex AI search tool
|
||||
assert llm_request.config.tools is not None
|
||||
assert len(llm_request.config.tools) == 2
|
||||
assert llm_request.config.tools[0] == existing_tool
|
||||
assert llm_request.config.tools[1].retrieval is not None
|
||||
assert llm_request.config.tools[1].retrieval.vertex_ai_search is not None
|
||||
@@ -0,0 +1,284 @@
|
||||
# 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.
|
||||
|
||||
"""Tests for model name utility functions."""
|
||||
|
||||
from google.adk.utils.model_name_utils import extract_model_name
|
||||
from google.adk.utils.model_name_utils import is_gemini_1_model
|
||||
from google.adk.utils.model_name_utils import is_gemini_2_model
|
||||
from google.adk.utils.model_name_utils import is_gemini_model
|
||||
|
||||
|
||||
class TestExtractModelName:
|
||||
"""Test the extract_model_name function."""
|
||||
|
||||
def test_extract_model_name_simple_model(self):
|
||||
"""Test extraction of simple model names."""
|
||||
assert extract_model_name('gemini-2.5-pro') == 'gemini-2.5-pro'
|
||||
assert extract_model_name('gemini-1.5-flash') == 'gemini-1.5-flash'
|
||||
assert extract_model_name('gemini-1.0-pro') == 'gemini-1.0-pro'
|
||||
assert extract_model_name('claude-3-sonnet') == 'claude-3-sonnet'
|
||||
assert extract_model_name('gpt-4') == 'gpt-4'
|
||||
|
||||
def test_extract_model_name_path_based_model(self):
|
||||
"""Test extraction of path-based model names."""
|
||||
path_model = 'projects/265104255505/locations/us-central1/publishers/google/models/gemini-2.0-flash-001'
|
||||
assert extract_model_name(path_model) == 'gemini-2.0-flash-001'
|
||||
|
||||
path_model_2 = 'projects/12345/locations/us-east1/publishers/google/models/gemini-1.5-pro-preview'
|
||||
assert extract_model_name(path_model_2) == 'gemini-1.5-pro-preview'
|
||||
|
||||
path_model_3 = 'projects/test-project/locations/europe-west1/publishers/google/models/claude-3-sonnet'
|
||||
assert extract_model_name(path_model_3) == 'claude-3-sonnet'
|
||||
|
||||
def test_extract_model_name_invalid_path(self):
|
||||
"""Test that invalid path formats return the original string."""
|
||||
invalid_paths = [
|
||||
'projects/invalid/path/format',
|
||||
'invalid/path/format',
|
||||
'projects/123/locations/us-central1/models/gemini-2.0-flash', # missing publishers
|
||||
'projects/123/publishers/google/models/gemini-2.0-flash', # missing locations
|
||||
'projects/123/locations/us-central1/publishers/google/gemini-2.0-flash', # missing models
|
||||
]
|
||||
|
||||
for invalid_path in invalid_paths:
|
||||
assert extract_model_name(invalid_path) == invalid_path
|
||||
|
||||
def test_extract_model_name_empty_string(self):
|
||||
"""Test extraction from empty string."""
|
||||
assert extract_model_name('') == ''
|
||||
|
||||
def test_extract_model_name_edge_cases(self):
|
||||
"""Test edge cases for model name extraction."""
|
||||
# Test with unusual but valid path patterns
|
||||
path_with_numbers = 'projects/123456789/locations/us-central1/publishers/google/models/gemini-2.0-flash-001'
|
||||
assert extract_model_name(path_with_numbers) == 'gemini-2.0-flash-001'
|
||||
|
||||
# Test with hyphens in project/location names
|
||||
path_with_hyphens = 'projects/my-test-project/locations/us-central1/publishers/google/models/gemini-1.5-pro'
|
||||
assert extract_model_name(path_with_hyphens) == 'gemini-1.5-pro'
|
||||
|
||||
|
||||
class TestIsGeminiModel:
|
||||
"""Test the is_gemini_model function."""
|
||||
|
||||
def test_is_gemini_model_simple_names(self):
|
||||
"""Test Gemini model detection with simple model names."""
|
||||
assert is_gemini_model('gemini-2.5-pro') is True
|
||||
assert is_gemini_model('gemini-1.5-flash') is True
|
||||
assert is_gemini_model('gemini-1.0-pro') is True
|
||||
assert is_gemini_model('gemini-2.0-flash-001') is True
|
||||
assert is_gemini_model('claude-3-sonnet') is False
|
||||
assert is_gemini_model('gpt-4') is False
|
||||
assert is_gemini_model('llama-2') is False
|
||||
|
||||
def test_is_gemini_model_path_based_names(self):
|
||||
"""Test Gemini model detection with path-based model names."""
|
||||
gemini_path = 'projects/265104255505/locations/us-central1/publishers/google/models/gemini-2.0-flash-001'
|
||||
assert is_gemini_model(gemini_path) is True
|
||||
|
||||
gemini_path_2 = 'projects/12345/locations/us-east1/publishers/google/models/gemini-1.5-pro-preview'
|
||||
assert is_gemini_model(gemini_path_2) is True
|
||||
|
||||
non_gemini_path = 'projects/265104255505/locations/us-central1/publishers/google/models/claude-3-sonnet'
|
||||
assert is_gemini_model(non_gemini_path) is False
|
||||
|
||||
def test_is_gemini_model_edge_cases(self):
|
||||
"""Test edge cases for Gemini model detection."""
|
||||
# Test with None
|
||||
assert is_gemini_model(None) is False
|
||||
|
||||
# Test with empty string
|
||||
assert is_gemini_model('') is False
|
||||
|
||||
# Test with model names containing gemini but not starting with it
|
||||
assert is_gemini_model('my-gemini-model') is False
|
||||
assert is_gemini_model('claude-gemini-hybrid') is False
|
||||
|
||||
# Test with model names that have gemini in the middle of the path
|
||||
tricky_path = 'projects/265104255505/locations/us-central1/publishers/gemini/models/claude-3-sonnet'
|
||||
assert is_gemini_model(tricky_path) is False
|
||||
|
||||
# Test with just "gemini" without dash
|
||||
assert is_gemini_model('gemini') is False
|
||||
assert is_gemini_model('gemini_1_5_flash') is False
|
||||
|
||||
def test_is_gemini_model_case_sensitivity(self):
|
||||
"""Test that model detection is case sensitive."""
|
||||
assert is_gemini_model('Gemini-2.5-pro') is False
|
||||
assert is_gemini_model('GEMINI-2.5-pro') is False
|
||||
assert is_gemini_model('gemini-2.5-PRO') is True # Only the start matters
|
||||
|
||||
|
||||
class TestIsGemini1Model:
|
||||
"""Test the is_gemini_1_model function."""
|
||||
|
||||
def test_is_gemini_1_model_simple_names(self):
|
||||
"""Test Gemini 1.x model detection with simple model names."""
|
||||
assert is_gemini_1_model('gemini-1.5-flash') is True
|
||||
assert is_gemini_1_model('gemini-1.0-pro') is True
|
||||
assert is_gemini_1_model('gemini-1.5-pro-preview') is True
|
||||
assert is_gemini_1_model('gemini-1.9-experimental') is True
|
||||
assert is_gemini_1_model('gemini-2.0-flash') is False
|
||||
assert is_gemini_1_model('gemini-2.5-pro') is False
|
||||
assert is_gemini_1_model('gemini-10.0-pro') is False # Only 1.x versions
|
||||
assert is_gemini_1_model('claude-3-sonnet') is False
|
||||
|
||||
def test_is_gemini_1_model_path_based_names(self):
|
||||
"""Test Gemini 1.x model detection with path-based model names."""
|
||||
gemini_1_path = 'projects/265104255505/locations/us-central1/publishers/google/models/gemini-1.5-flash-001'
|
||||
assert is_gemini_1_model(gemini_1_path) is True
|
||||
|
||||
gemini_1_path_2 = 'projects/12345/locations/us-east1/publishers/google/models/gemini-1.0-pro-preview'
|
||||
assert is_gemini_1_model(gemini_1_path_2) is True
|
||||
|
||||
gemini_2_path = 'projects/265104255505/locations/us-central1/publishers/google/models/gemini-2.0-flash-001'
|
||||
assert is_gemini_1_model(gemini_2_path) is False
|
||||
|
||||
def test_is_gemini_1_model_edge_cases(self):
|
||||
"""Test edge cases for Gemini 1.x model detection."""
|
||||
# Test with None
|
||||
assert is_gemini_1_model(None) is False
|
||||
|
||||
# Test with empty string
|
||||
assert is_gemini_1_model('') is False
|
||||
|
||||
# Test with model names containing gemini-1 but not starting with it
|
||||
assert is_gemini_1_model('my-gemini-1.5-model') is False
|
||||
assert is_gemini_1_model('custom-gemini-1.5-flash') is False
|
||||
|
||||
# Test with invalid versions
|
||||
assert is_gemini_1_model('gemini-1') is False # Missing dot
|
||||
assert is_gemini_1_model('gemini-1-pro') is False # Missing dot
|
||||
assert is_gemini_1_model('gemini-1.') is False # Missing version number
|
||||
|
||||
|
||||
class TestIsGemini2Model:
|
||||
"""Test the is_gemini_2_model function."""
|
||||
|
||||
def test_is_gemini_2_model_simple_names(self):
|
||||
"""Test Gemini 2.x model detection with simple model names."""
|
||||
assert is_gemini_2_model('gemini-2.0-flash') is True
|
||||
assert is_gemini_2_model('gemini-2.5-pro') is True
|
||||
assert is_gemini_2_model('gemini-2.0-flash-001') is True
|
||||
assert is_gemini_2_model('gemini-2.9-experimental') is True
|
||||
assert is_gemini_2_model('gemini-1.5-flash') is False
|
||||
assert is_gemini_2_model('gemini-1.0-pro') is False
|
||||
assert is_gemini_2_model('gemini-3.0-pro') is False # Only 2.x versions
|
||||
assert is_gemini_2_model('claude-3-sonnet') is False
|
||||
|
||||
def test_is_gemini_2_model_path_based_names(self):
|
||||
"""Test Gemini 2.x model detection with path-based model names."""
|
||||
gemini_2_path = 'projects/265104255505/locations/us-central1/publishers/google/models/gemini-2.0-flash-001'
|
||||
assert is_gemini_2_model(gemini_2_path) is True
|
||||
|
||||
gemini_2_path_2 = 'projects/12345/locations/us-east1/publishers/google/models/gemini-2.5-pro-preview'
|
||||
assert is_gemini_2_model(gemini_2_path_2) is True
|
||||
|
||||
gemini_1_path = 'projects/265104255505/locations/us-central1/publishers/google/models/gemini-1.5-flash-001'
|
||||
assert is_gemini_2_model(gemini_1_path) is False
|
||||
|
||||
def test_is_gemini_2_model_edge_cases(self):
|
||||
"""Test edge cases for Gemini 2.x model detection."""
|
||||
# Test with None
|
||||
assert is_gemini_2_model(None) is False
|
||||
|
||||
# Test with empty string
|
||||
assert is_gemini_2_model('') is False
|
||||
|
||||
# Test with model names containing gemini-2 but not starting with it
|
||||
assert is_gemini_2_model('my-gemini-2.5-model') is False
|
||||
assert is_gemini_2_model('custom-gemini-2.0-flash') is False
|
||||
|
||||
# Test with invalid versions
|
||||
assert is_gemini_2_model('gemini-2') is False # Missing dot
|
||||
assert is_gemini_2_model('gemini-2-pro') is False # Missing dot
|
||||
assert is_gemini_2_model('gemini-2.') is False # Missing version number
|
||||
|
||||
|
||||
class TestModelNameUtilsIntegration:
|
||||
"""Integration tests for model name utilities."""
|
||||
|
||||
def test_model_classification_consistency(self):
|
||||
"""Test that model classification functions are consistent."""
|
||||
test_models = [
|
||||
'gemini-1.5-flash',
|
||||
'gemini-2.0-flash',
|
||||
'gemini-2.5-pro',
|
||||
'projects/123/locations/us-central1/publishers/google/models/gemini-1.5-pro',
|
||||
'projects/123/locations/us-central1/publishers/google/models/gemini-2.0-flash',
|
||||
'claude-3-sonnet',
|
||||
'gpt-4',
|
||||
]
|
||||
|
||||
for model in test_models:
|
||||
# A model can only be either Gemini 1.x or Gemini 2.x, not both
|
||||
if is_gemini_1_model(model):
|
||||
assert not is_gemini_2_model(
|
||||
model
|
||||
), f'Model {model} classified as both Gemini 1.x and 2.x'
|
||||
assert is_gemini_model(
|
||||
model
|
||||
), f'Model {model} is Gemini 1.x but not classified as Gemini'
|
||||
|
||||
if is_gemini_2_model(model):
|
||||
assert not is_gemini_1_model(
|
||||
model
|
||||
), f'Model {model} classified as both Gemini 1.x and 2.x'
|
||||
assert is_gemini_model(
|
||||
model
|
||||
), f'Model {model} is Gemini 2.x but not classified as Gemini'
|
||||
|
||||
# If it's neither Gemini 1.x nor 2.x, it should not be classified as Gemini
|
||||
if not is_gemini_1_model(model) and not is_gemini_2_model(model):
|
||||
if model and 'gemini-' not in extract_model_name(model):
|
||||
assert not is_gemini_model(
|
||||
model
|
||||
), f'Non-Gemini model {model} classified as Gemini'
|
||||
|
||||
def test_path_vs_simple_model_consistency(self):
|
||||
"""Test that path-based and simple model names are classified consistently."""
|
||||
model_pairs = [
|
||||
(
|
||||
'gemini-1.5-flash',
|
||||
'projects/123/locations/us-central1/publishers/google/models/gemini-1.5-flash',
|
||||
),
|
||||
(
|
||||
'gemini-2.0-flash',
|
||||
'projects/123/locations/us-central1/publishers/google/models/gemini-2.0-flash',
|
||||
),
|
||||
(
|
||||
'gemini-2.5-pro',
|
||||
'projects/123/locations/us-central1/publishers/google/models/gemini-2.5-pro',
|
||||
),
|
||||
(
|
||||
'claude-3-sonnet',
|
||||
'projects/123/locations/us-central1/publishers/google/models/claude-3-sonnet',
|
||||
),
|
||||
]
|
||||
|
||||
for simple_model, path_model in model_pairs:
|
||||
# Both forms should be classified identically
|
||||
assert is_gemini_model(simple_model) == is_gemini_model(path_model), (
|
||||
f'Inconsistent Gemini classification for {simple_model} vs'
|
||||
f' {path_model}'
|
||||
)
|
||||
assert is_gemini_1_model(simple_model) == is_gemini_1_model(path_model), (
|
||||
f'Inconsistent Gemini 1.x classification for {simple_model} vs'
|
||||
f' {path_model}'
|
||||
)
|
||||
assert is_gemini_2_model(simple_model) == is_gemini_2_model(path_model), (
|
||||
f'Inconsistent Gemini 2.x classification for {simple_model} vs'
|
||||
f' {path_model}'
|
||||
)
|
||||
Reference in New Issue
Block a user