mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
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:
committed by
Copybara-Service
parent
6da7274858
commit
e63180cb62
@@ -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",
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user