feat: Add ADK version check and set MemoryBankService as default when google-adk>=1.5.0

PiperOrigin-RevId: 781290630

Author: vertex-sdk-bot

tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py

@@ -84,10 +84,10 @@ def simple_span_processor_mock():
 
 
 @pytest.fixture
-def mock_adk_major_version():
+def mock_adk_version():
     with mock.patch(
-        "google.cloud.aiplatform.vertexai.preview.reasoning_engines.templates.adk.get_adk_major_version",
-        return_value=1,
+        "google.cloud.aiplatform.vertexai.preview.reasoning_engines.templates.adk.get_adk_version",
+        return_value="1.5.0",
     ):
         yield
 
@@ -148,17 +148,17 @@ async def run_async(self, *args, **kwargs):
         )
 
 
-@pytest.mark.usefixtures("google_auth_mock", "mock_adk_major_version")
+@pytest.mark.usefixtures("google_auth_mock")
 class TestAdkApp:
-    def test_adk_major_version(self):
+    def test_adk_version(self):
         with mock.patch(
-            "google.cloud.aiplatform.vertexai.preview.reasoning_engines.templates.adk.get_adk_major_version",
-            return_value=0,
+            "google.cloud.aiplatform.vertexai.preview.reasoning_engines.templates.adk.get_adk_version",
+            return_value="0.5.0",
         ):
             with pytest.raises(
                 ValueError,
                 match=(
-                    "Unsupported google-adk major version: 0, please use"
+                    "Unsupported google-adk version: 0.5.0, please use"
                     " google-adk>=1.0.0 for AdkApp deployment."
                 ),
             ):
@@ -400,7 +400,7 @@ def test_enable_tracing_warning(self, caplog):
         # assert "enable_tracing=True but proceeding with tracing disabled" in caplog.text
 
 
-@pytest.mark.usefixtures("mock_adk_major_version")
+@pytest.mark.usefixtures("mock_adk_version")
 class TestAdkAppErrors:
     def test_raise_get_session_not_found_error(self):
         with pytest.raises(

vertexai/preview/reasoning_engines/templates/adk.py

@@ -81,14 +81,31 @@
 _DEFAULT_USER_ID = "default-user-id"
 
 
-def get_adk_major_version() -> int:
-    """Get the major version of google-adk."""
+def get_adk_version() -> Optional[str]:
+    """Returns the version of the ADK package."""
     try:
         from google.adk import version
 
-        return int(version.__version__.split(".")[0])
+        return version.__version__
     except ImportError:
-        return 0
+        return None
+
+
+def is_version_sufficient(version_to_check: str) -> bool:
+    """Compares the existing version of ADK with the required version.
+
+    Args:
+        version_to_check: The version string to check.
+
+    Returns:
+        True if the existing version is sufficient, False otherwise.
+    """
+    try:
+        from packaging.version import parse
+
+        return parse(get_adk_version()) >= parse(version_to_check)
+    except (AttributeError, ImportError):
+        return False
 
 
 class _ArtifactVersion:
@@ -294,10 +311,10 @@ def __init__(
         """An ADK Application."""
         from google.cloud.aiplatform import initializer
 
-        adk_major_version = get_adk_major_version()
-        if adk_major_version < 1:
+        adk_version = get_adk_version()
+        if not is_version_sufficient("1.0.0"):
             msg = (
-                f"Unsupported google-adk major version: {adk_major_version}, "
+                f"Unsupported google-adk version: {adk_version}, "
                 "please use google-adk>=1.0.0 for AdkApp deployment."
             )
             raise ValueError(msg)
@@ -464,16 +481,35 @@ def set_up(self):
                 VertexAiSessionService,
             )
 
-            self._tmpl_attrs["session_service"] = VertexAiSessionService(
-                project=project,
-                location=location,
-            )
+            if is_version_sufficient("1.5.0"):
+                self._tmpl_attrs["session_service"] = VertexAiSessionService(
+                    project=project,
+                    location=location,
+                    agent_engine_id=os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID"),
+                )
+            else:
+                self._tmpl_attrs["session_service"] = VertexAiSessionService(
+                    project=project,
+                    location=location,
+                )
         else:
             self._tmpl_attrs["session_service"] = InMemorySessionService()
 
         memory_service_builder = self._tmpl_attrs.get("memory_service_builder")
         if memory_service_builder:
             self._tmpl_attrs["memory_service"] = memory_service_builder()
+        elif "GOOGLE_CLOUD_AGENT_ENGINE_ID" in os.environ and is_version_sufficient(
+            "1.5.0"
+        ):
+            from google.adk.memory.vertex_ai_memory_bank_service import (
+                VertexAiMemoryBankService,
+            )
+
+            self._tmpl_attrs["memory_service"] = VertexAiMemoryBankService(
+                project=project,
+                location=location,
+                agent_engine_id=os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID"),
+            )
         else:
             self._tmpl_attrs["memory_service"] = InMemoryMemoryService()