feat: Add the support for returning struct_data.uri in DiscoveryEngineSearchTool

For https://github.com/google/adk-python/issues/3146

PiperOrigin-RevId: 818458080
This commit is contained in:
Xuan Yang
2025-10-12 20:28:08 -07:00
committed by Copybara-Service
parent 6da7274858
commit e63180cb62
2 changed files with 19 additions and 5 deletions
@@ -80,7 +80,7 @@ class DiscoveryEngineSearchTool(FunctionTool):
self,
query: str,
) -> dict[str, Any]:
"""Search the discovery engine.
"""Search through Vertex AI Search's discovery engine search API.
Args:
query: The search query.
@@ -113,12 +113,22 @@ class DiscoveryEngineSearchTool(FunctionTool):
response = self._discovery_engine_client.search(request)
for item in response.results:
chunk = item.chunk
if not chunk or not chunk.document_metadata:
if not chunk:
continue
title = ""
uri = ""
doc_metadata = chunk.document_metadata
if doc_metadata:
title = doc_metadata.title
uri = doc_metadata.uri
# Prioritize URI from struct_data if it exists.
if doc_metadata.struct_data and "uri" in doc_metadata.struct_data:
uri = doc_metadata.struct_data["uri"]
results.append({
"title": chunk.document_metadata.title,
"url": chunk.document_metadata.uri,
"title": title,
"url": uri,
"content": chunk.content,
})
except GoogleAPICallError as e:
@@ -87,7 +87,11 @@ class TestDiscoveryEngineSearchTool:
chunk=discoveryengine.Chunk(
document_metadata={
"title": "Test Title",
"uri": "http://example.com",
"uri": "gs://test_bucket/test_file",
"struct_data": {
"key1": "value1",
"uri": "http://example.com",
},
},
content="Test Content",
)