From 9478a31bf2257f0b668ae7eb91a10863e87c7bed Mon Sep 17 00:00:00 2001
From: "Xiang (Sean) Zhou" <seanzhougoogle@google.com>
Date: Fri, 8 Aug 2025 11:37:12 -0700
Subject: [PATCH] fix: Lazy load retrieval tools and prompt users to install
 extensions if import failed

PiperOrigin-RevId: 792701550
---
 src/google/adk/tools/retrieval/__init__.py | 42 +++++++++++++++++-----
 1 file changed, 33 insertions(+), 9 deletions(-)

diff --git a/src/google/adk/tools/retrieval/__init__.py b/src/google/adk/tools/retrieval/__init__.py
index 53778061..f5495d4d 100644
--- a/src/google/adk/tools/retrieval/__init__.py
+++ b/src/google/adk/tools/retrieval/__init__.py
@@ -13,20 +13,44 @@
 # limitations under the License.
 
 from .base_retrieval_tool import BaseRetrievalTool
-from .files_retrieval import FilesRetrieval
-from .llama_index_retrieval import LlamaIndexRetrieval
 
 __all__ = [
-    'BaseRetrievalTool',
-    'FilesRetrieval',
-    'LlamaIndexRetrieval',
-    'VertexAiRagRetrieval',
+    "BaseRetrievalTool",
+    "FilesRetrieval",
+    "LlamaIndexRetrieval",
+    "VertexAiRagRetrieval",
 ]
 
 
 def __getattr__(name: str):
-  if name == 'VertexAiRagRetrieval':
-    from .vertex_ai_rag_retrieval import VertexAiRagRetrieval
+  if name == "FilesRetrieval":
+    try:
+      from .files_retrieval import FilesRetrieval
 
-    return VertexAiRagRetrieval
+      return FilesRetrieval
+    except ImportError as e:
+      raise ImportError(
+          "FilesRetrieval requires additional dependencies. "
+          'Please install with: pip install "google-adk[extensions]"'
+      ) from e
+  elif name == "LlamaIndexRetrieval":
+    try:
+      from .llama_index_retrieval import LlamaIndexRetrieval
+
+      return LlamaIndexRetrieval
+    except ImportError as e:
+      raise ImportError(
+          "LlamaIndexRetrieval requires additional dependencies. "
+          'Please install with: pip install "google-adk[extensions]"'
+      ) from e
+  elif name == "VertexAiRagRetrieval":
+    try:
+      from .vertex_ai_rag_retrieval import VertexAiRagRetrieval
+
+      return VertexAiRagRetrieval
+    except ImportError as e:
+      raise ImportError(
+          "VertexAiRagRetrieval requires additional dependencies. "
+          'Please install with: pip install "google-adk[extensions]"'
+      ) from e
   raise AttributeError(f"module '{__name__}' has no attribute '{name}'")