feat: Use json schema for base_retrieval_tool, load_artifacts_tool, and load_memory_tool declaration when feature enabled

Co-authored-by: Xuan Yang <xygoogle@google.com>
PiperOrigin-RevId: 858435881

src/google/adk/tools/load_artifacts_tool.py

@@ -24,6 +24,8 @@ from typing import TYPE_CHECKING
 from google.genai import types
 from typing_extensions import override
 
+from ..features import FeatureName
+from ..features import is_feature_enabled
 from .base_tool import BaseTool
 
 # MIME types Gemini accepts for inline data in requests.
@@ -132,6 +134,20 @@ web UI)."""),
     )
 
   def _get_declaration(self) -> types.FunctionDeclaration | None:
+    if is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL):
+      return types.FunctionDeclaration(
+          name=self.name,
+          description=self.description,
+          parameters_json_schema={
+              'type': 'object',
+              'properties': {
+                  'artifact_names': {
+                      'type': 'array',
+                      'items': {'type': 'string'},
+                  },
+              },
+          },
+      )
     return types.FunctionDeclaration(
         name=self.name,
         description=self.description,

src/google/adk/tools/load_memory_tool.py

@@ -21,6 +21,8 @@ from pydantic import BaseModel
 from pydantic import Field
 from typing_extensions import override
 
+from ..features import FeatureName
+from ..features import is_feature_enabled
 from ..memory.memory_entry import MemoryEntry
 from .function_tool import FunctionTool
 from .tool_context import ToolContext
@@ -59,6 +61,18 @@ class LoadMemoryTool(FunctionTool):
 
   @override
   def _get_declaration(self) -> types.FunctionDeclaration | None:
+    if is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL):
+      return types.FunctionDeclaration(
+          name=self.name,
+          description=self.description,
+          parameters_json_schema={
+              'type': 'object',
+              'properties': {
+                  'query': {'type': 'string'},
+              },
+              'required': ['query'],
+          },
+      )
     return types.FunctionDeclaration(
         name=self.name,
         description=self.description,

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

@@ -12,9 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 from google.genai import types
 from typing_extensions import override
 
+from ...features import FeatureName
+from ...features import is_feature_enabled
 from ..base_tool import BaseTool
 
 
@@ -22,6 +26,20 @@ class BaseRetrievalTool(BaseTool):
 
   @override
   def _get_declaration(self) -> types.FunctionDeclaration:
+    if is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL):
+      return types.FunctionDeclaration(
+          name=self.name,
+          description=self.description,
+          parameters_json_schema={
+              'type': 'object',
+              'properties': {
+                  'query': {
+                      'type': 'string',
+                      'description': 'The query to retrieve.',
+                  },
+              },
+          },
+      )
     return types.FunctionDeclaration(
         name=self.name,
         description=self.description,