Skip to content

Commit 9478a31

Browse files
seanzhougooglecopybara-github
authored andcommitted
fix: Lazy load retrieval tools and prompt users to install extensions if import failed
PiperOrigin-RevId: 792701550
1 parent f2005a2 commit 9478a31

File tree

1 file changed

+33
-9
lines changed

1 file changed

+33
-9
lines changed

src/google/adk/tools/retrieval/__init__.py

Lines changed: 33 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -13,20 +13,44 @@
1313
# limitations under the License.
1414

1515
from .base_retrieval_tool import BaseRetrievalTool
16-
from .files_retrieval import FilesRetrieval
17-
from .llama_index_retrieval import LlamaIndexRetrieval
1816

1917
__all__ = [
20-
'BaseRetrievalTool',
21-
'FilesRetrieval',
22-
'LlamaIndexRetrieval',
23-
'VertexAiRagRetrieval',
18+
"BaseRetrievalTool",
19+
"FilesRetrieval",
20+
"LlamaIndexRetrieval",
21+
"VertexAiRagRetrieval",
2422
]
2523

2624

2725
def __getattr__(name: str):
28-
if name == 'VertexAiRagRetrieval':
29-
from .vertex_ai_rag_retrieval import VertexAiRagRetrieval
26+
if name == "FilesRetrieval":
27+
try:
28+
from .files_retrieval import FilesRetrieval
3029

31-
return VertexAiRagRetrieval
30+
return FilesRetrieval
31+
except ImportError as e:
32+
raise ImportError(
33+
"FilesRetrieval requires additional dependencies. "
34+
'Please install with: pip install "google-adk[extensions]"'
35+
) from e
36+
elif name == "LlamaIndexRetrieval":
37+
try:
38+
from .llama_index_retrieval import LlamaIndexRetrieval
39+
40+
return LlamaIndexRetrieval
41+
except ImportError as e:
42+
raise ImportError(
43+
"LlamaIndexRetrieval requires additional dependencies. "
44+
'Please install with: pip install "google-adk[extensions]"'
45+
) from e
46+
elif name == "VertexAiRagRetrieval":
47+
try:
48+
from .vertex_ai_rag_retrieval import VertexAiRagRetrieval
49+
50+
return VertexAiRagRetrieval
51+
except ImportError as e:
52+
raise ImportError(
53+
"VertexAiRagRetrieval requires additional dependencies. "
54+
'Please install with: pip install "google-adk[extensions]"'
55+
) from e
3256
raise AttributeError(f"module '{__name__}' has no attribute '{name}'")

0 commit comments

Comments
 (0)