feat: Support new embedding model in files retrieval

Co-authored-by: Xiang (Sean) Zhou <seanzhougoogle@google.com>
PiperOrigin-RevId: 883401159

Author: seanzhougoogle

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

@@ -32,15 +32,18 @@ def _get_default_embedding_model() -> BaseEmbedding:
   """Get the default Google Gemini embedding model.
 
   Returns:
-    GoogleGenAIEmbedding instance configured with text-embedding-004 model.
+    GoogleGenAIEmbedding instance configured with gemini-embedding-2-preview model.
 
   Raises:
     ImportError: If llama-index-embeddings-google-genai package is not installed.
   """
   try:
     from llama_index.embeddings.google_genai import GoogleGenAIEmbedding
 
-    return GoogleGenAIEmbedding(model_name="text-embedding-004")
+    return GoogleGenAIEmbedding(
+        model_name="gemini-embedding-2-preview",
+        embed_batch_size=1,
+    )
   except ImportError as e:
     raise ImportError(
         "llama-index-embeddings-google-genai package not found. "
@@ -65,7 +68,7 @@ def __init__(
       description: Description of the tool.
       input_dir: Directory path containing files to index.
       embedding_model: Optional custom embedding model. If None, defaults to
-        Google's text-embedding-004 model.
+        Google's gemini-embedding-2-preview model.
     """
     self.input_dir = input_dir
 

tests/unittests/tools/retrieval/test_files_retrieval.py

@@ -122,7 +122,8 @@ def test_get_default_embedding_model_success(self):
       result = _get_default_embedding_model()
 
       mock_module.GoogleGenAIEmbedding.assert_called_once_with(
-          model_name="text-embedding-004"
+          model_name="gemini-embedding-2-preview",
+          embed_batch_size=1,
       )
       assert result == mock_embedding_instance