feat(tools): Add debug logging to VertexAiSearchTool

Merge https://github.com/google/adk-python/pull/3284

**Problem:**
When debugging agents that utilize the `VertexAiSearchTool`, it's currently difficult to inspect the specific configuration parameters (datastore ID, engine ID, filter, max_results, etc.) being passed to the underlying Vertex AI Search API via the `LlmRequest`. This lack of visibility can hinder troubleshooting efforts related to tool configuration.

**Solution:**
This PR enhances the `VertexAiSearchTool` by adding a **debug-level log statement** within the `process_llm_request` method. This log precisely records the parameters being used for the Vertex AI Search configuration just before it's appended to the `LlmRequest`.

This provides developers with crucial visibility into the tool's runtime behavior when debug logging is enabled, significantly improving the **debuggability** of agents using this tool. Corresponding unit tests were updated to rigorously verify this new logging output using `caplog`. Additionally, minor fixes were made to the tests to resolve Pydantic validation errors.

Co-authored-by: Xuan Yang <xygoogle@google.com>
COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/3284 from omkute10:feat/add-logging-vertex-search-tool 199c12bf00a57abe202401591088c0423b39b928
PiperOrigin-RevId: 836419886
This commit is contained in:
Om Kute
2025-11-24 17:29:50 -08:00
committed by Copybara-Service
parent b331d97dfb
commit c6e7d6b16a
2 changed files with 170 additions and 12 deletions
@@ -14,6 +14,7 @@
from __future__ import annotations from __future__ import annotations
import logging
from typing import Optional from typing import Optional
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
@@ -25,6 +26,8 @@ from ..utils.model_name_utils import is_gemini_model
from .base_tool import BaseTool from .base_tool import BaseTool
from .tool_context import ToolContext from .tool_context import ToolContext
logger = logging.getLogger('google_adk.' + __name__)
if TYPE_CHECKING: if TYPE_CHECKING:
from ..models import LlmRequest from ..models import LlmRequest
@@ -102,6 +105,30 @@ class VertexAiSearchTool(BaseTool):
) )
llm_request.config = llm_request.config or types.GenerateContentConfig() llm_request.config = llm_request.config or types.GenerateContentConfig()
llm_request.config.tools = llm_request.config.tools or [] llm_request.config.tools = llm_request.config.tools or []
# Format data_store_specs concisely for logging
if self.data_store_specs:
spec_ids = [
spec.data_store.split('/')[-1] if spec.data_store else 'unnamed'
for spec in self.data_store_specs
]
specs_info = (
f'{len(self.data_store_specs)} spec(s): [{", ".join(spec_ids)}]'
)
else:
specs_info = None
logger.debug(
'Adding Vertex AI Search tool config to LLM request: '
'datastore=%s, engine=%s, filter=%s, max_results=%s, '
'data_store_specs=%s',
self.data_store_id,
self.search_engine_id,
self.filter,
self.max_results,
specs_info,
)
llm_request.config.tools.append( llm_request.config.tools.append(
types.Tool( types.Tool(
retrieval=types.Retrieval( retrieval=types.Retrieval(
@@ -12,6 +12,8 @@
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
import logging
from google.adk.agents.invocation_context import InvocationContext from google.adk.agents.invocation_context import InvocationContext
from google.adk.agents.sequential_agent import SequentialAgent from google.adk.agents.sequential_agent import SequentialAgent
from google.adk.models.llm_request import LlmRequest from google.adk.models.llm_request import LlmRequest
@@ -24,6 +26,10 @@ from google.adk.utils.model_name_utils import is_gemini_model
from google.genai import types from google.genai import types
import pytest import pytest
VERTEX_SEARCH_TOOL_LOGGER_NAME = (
'google_adk.google.adk.tools.vertex_ai_search_tool'
)
async def _create_tool_context() -> ToolContext: async def _create_tool_context() -> ToolContext:
session_service = InMemorySessionService() session_service = InMemorySessionService()
@@ -121,12 +127,34 @@ class TestVertexAiSearchTool:
tool = VertexAiSearchTool(data_store_id='test_data_store') tool = VertexAiSearchTool(data_store_id='test_data_store')
assert tool.data_store_id == 'test_data_store' assert tool.data_store_id == 'test_data_store'
assert tool.search_engine_id is None assert tool.search_engine_id is None
assert tool.data_store_specs is None
def test_init_with_search_engine_id(self): def test_init_with_search_engine_id(self):
"""Test initialization with search engine ID.""" """Test initialization with search engine ID."""
tool = VertexAiSearchTool(search_engine_id='test_search_engine') tool = VertexAiSearchTool(search_engine_id='test_search_engine')
assert tool.search_engine_id == 'test_search_engine' assert tool.search_engine_id == 'test_search_engine'
assert tool.data_store_id is None assert tool.data_store_id is None
assert tool.data_store_specs is None
def test_init_with_engine_and_specs(self):
"""Test initialization with search engine ID and specs."""
specs = [
types.VertexAISearchDataStoreSpec(
dataStore=(
'projects/p/locations/l/collections/c/dataStores/spec_store'
)
)
]
engine_id = (
'projects/p/locations/l/collections/c/engines/test_search_engine'
)
tool = VertexAiSearchTool(
search_engine_id=engine_id,
data_store_specs=specs,
)
assert tool.search_engine_id == engine_id
assert tool.data_store_id is None
assert tool.data_store_specs == specs
def test_init_with_neither_raises_error(self): def test_init_with_neither_raises_error(self):
"""Test that initialization without either ID raises ValueError.""" """Test that initialization without either ID raises ValueError."""
@@ -146,10 +174,34 @@ class TestVertexAiSearchTool:
data_store_id='test_data_store', search_engine_id='test_search_engine' data_store_id='test_data_store', search_engine_id='test_search_engine'
) )
def test_init_with_specs_but_no_engine_raises_error(self):
"""Test that specs without engine ID raises ValueError."""
specs = [
types.VertexAISearchDataStoreSpec(
dataStore=(
'projects/p/locations/l/collections/c/dataStores/spec_store'
)
)
]
with pytest.raises(
ValueError,
match=(
'search_engine_id must be specified if data_store_specs is'
' specified'
),
):
VertexAiSearchTool(
data_store_id='test_data_store', data_store_specs=specs
)
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_process_llm_request_with_simple_gemini_model(self): async def test_process_llm_request_with_simple_gemini_model(self, caplog):
"""Test processing LLM request with simple Gemini model name.""" """Test processing LLM request with simple Gemini model name."""
tool = VertexAiSearchTool(data_store_id='test_data_store') caplog.set_level(logging.DEBUG, logger=VERTEX_SEARCH_TOOL_LOGGER_NAME)
tool = VertexAiSearchTool(
data_store_id='test_data_store', filter='f', max_results=5
)
tool_context = await _create_tool_context() tool_context = await _create_tool_context()
llm_request = LlmRequest( llm_request = LlmRequest(
@@ -162,17 +214,56 @@ class TestVertexAiSearchTool:
assert llm_request.config.tools is not None assert llm_request.config.tools is not None
assert len(llm_request.config.tools) == 1 assert len(llm_request.config.tools) == 1
assert llm_request.config.tools[0].retrieval is not None retrieval_tool = llm_request.config.tools[0]
assert llm_request.config.tools[0].retrieval.vertex_ai_search is not None assert retrieval_tool.retrieval is not None
assert retrieval_tool.retrieval.vertex_ai_search is not None
assert (
retrieval_tool.retrieval.vertex_ai_search.datastore == 'test_data_store'
)
assert retrieval_tool.retrieval.vertex_ai_search.engine is None
assert retrieval_tool.retrieval.vertex_ai_search.filter == 'f'
assert retrieval_tool.retrieval.vertex_ai_search.max_results == 5
# Verify debug log
debug_records = [
r
for r in caplog.records
if 'Adding Vertex AI Search tool config' in r.message
]
assert len(debug_records) == 1
log_message = debug_records[0].getMessage()
assert 'datastore=test_data_store' in log_message
assert 'engine=None' in log_message
assert 'filter=f' in log_message
assert 'max_results=5' in log_message
assert 'data_store_specs=None' in log_message
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_process_llm_request_with_path_based_gemini_model(self): async def test_process_llm_request_with_path_based_gemini_model(self, caplog):
"""Test processing LLM request with path-based Gemini model name.""" """Test processing LLM request with path-based Gemini model name."""
tool = VertexAiSearchTool(data_store_id='test_data_store') caplog.set_level(logging.DEBUG, logger=VERTEX_SEARCH_TOOL_LOGGER_NAME)
specs = [
types.VertexAISearchDataStoreSpec(
dataStore=(
'projects/p/locations/l/collections/c/dataStores/spec_store'
)
)
]
engine_id = 'projects/p/locations/l/collections/c/engines/test_engine'
tool = VertexAiSearchTool(
search_engine_id=engine_id,
data_store_specs=specs,
filter='f2',
max_results=10,
)
tool_context = await _create_tool_context() tool_context = await _create_tool_context()
llm_request = LlmRequest( llm_request = LlmRequest(
model='projects/265104255505/locations/us-central1/publishers/google/models/gemini-2.0-flash-001', model=(
'projects/265104255505/locations/us-central1/publishers/'
'google/models/gemini-2.0-flash-001'
),
config=types.GenerateContentConfig(), config=types.GenerateContentConfig(),
) )
@@ -182,8 +273,28 @@ class TestVertexAiSearchTool:
assert llm_request.config.tools is not None assert llm_request.config.tools is not None
assert len(llm_request.config.tools) == 1 assert len(llm_request.config.tools) == 1
assert llm_request.config.tools[0].retrieval is not None retrieval_tool = llm_request.config.tools[0]
assert llm_request.config.tools[0].retrieval.vertex_ai_search is not None assert retrieval_tool.retrieval is not None
assert retrieval_tool.retrieval.vertex_ai_search is not None
assert retrieval_tool.retrieval.vertex_ai_search.datastore is None
assert retrieval_tool.retrieval.vertex_ai_search.engine == engine_id
assert retrieval_tool.retrieval.vertex_ai_search.filter == 'f2'
assert retrieval_tool.retrieval.vertex_ai_search.max_results == 10
assert retrieval_tool.retrieval.vertex_ai_search.data_store_specs == specs
# Verify debug log
debug_records = [
r
for r in caplog.records
if 'Adding Vertex AI Search tool config' in r.message
]
assert len(debug_records) == 1
log_message = debug_records[0].getMessage()
assert 'datastore=None' in log_message
assert f'engine={engine_id}' in log_message
assert 'filter=f2' in log_message
assert 'max_results=10' in log_message
assert 'data_store_specs=1 spec(s): [spec_store]' in log_message
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_process_llm_request_with_gemini_1_and_other_tools_raises_error( async def test_process_llm_request_with_gemini_1_and_other_tools_raises_error(
@@ -291,9 +402,11 @@ class TestVertexAiSearchTool:
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_process_llm_request_with_gemini_2_and_other_tools_succeeds( async def test_process_llm_request_with_gemini_2_and_other_tools_succeeds(
self, self, caplog
): ):
"""Test that Gemini 2.x with other tools succeeds.""" """Test that Gemini 2.x with other tools succeeds."""
caplog.set_level(logging.DEBUG, logger=VERTEX_SEARCH_TOOL_LOGGER_NAME)
tool = VertexAiSearchTool(data_store_id='test_data_store') tool = VertexAiSearchTool(data_store_id='test_data_store')
tool_context = await _create_tool_context() tool_context = await _create_tool_context()
@@ -316,5 +429,23 @@ class TestVertexAiSearchTool:
assert llm_request.config.tools is not None assert llm_request.config.tools is not None
assert len(llm_request.config.tools) == 2 assert len(llm_request.config.tools) == 2
assert llm_request.config.tools[0] == existing_tool assert llm_request.config.tools[0] == existing_tool
assert llm_request.config.tools[1].retrieval is not None retrieval_tool = llm_request.config.tools[1]
assert llm_request.config.tools[1].retrieval.vertex_ai_search is not None assert retrieval_tool.retrieval is not None
assert retrieval_tool.retrieval.vertex_ai_search is not None
assert (
retrieval_tool.retrieval.vertex_ai_search.datastore == 'test_data_store'
)
# Verify debug log
debug_records = [
r
for r in caplog.records
if 'Adding Vertex AI Search tool config' in r.message
]
assert len(debug_records) == 1
log_message = debug_records[0].getMessage()
assert 'datastore=test_data_store' in log_message
assert 'engine=None' in log_message
assert 'filter=None' in log_message
assert 'max_results=None' in log_message
assert 'data_store_specs=None' in log_message