feat: Add new template for AgentEngine

PiperOrigin-RevId: 745364132

Author: yeesian

noxfile.py

@@ -212,6 +212,7 @@ def default(session):
         "--cov-report=",
         "--cov-fail-under=0",
         "--ignore=tests/unit/vertex_ray",
+        "--ignore=tests/unit/vertex_adk",
         "--ignore=tests/unit/vertex_langchain",
         "--ignore=tests/unit/vertex_ag2",
         "--ignore=tests/unit/vertex_llama_index",

setup.py

@@ -142,6 +142,10 @@
     "xgboost_ray",
 ]
 
+adk_extra_require = [
+    "google-adk >= 0.0.2",
+]
+
 reasoning_engine_extra_require = [
     "cloudpickle >= 3.0, < 4.0",
     "google-cloud-trace < 2",
@@ -320,6 +324,7 @@
         "preview": preview_extra_require,
         "ray": ray_extra_require,
         "ray_testing": ray_testing_extra_require,
+        "adk": adk_extra_require,
         "reasoningengine": reasoning_engine_extra_require,
         "agent_engines": agent_engines_extra_require,
         "evaluation": evaluation_extra_require,

testing/constraints-3.10.txt

@@ -15,4 +15,5 @@ ray==2.4.0 # Pinned until 2.9.3 is verified for Ray tests
 ipython==8.22.2 # Pinned to unbreak TypeAliasType import error
 scikit-learn!=1.4.1.post1 # Pin to unbreak test_sklearn (b/332610038)
 requests==2.31.0 # Pinned to unbreak http+docker error (b/342669351)
-google-vizier==0.1.21
+google-vizier==0.1.21
+google-adk==0.0.2

testing/constraints-3.11.txt

@@ -11,4 +11,5 @@ packaging==24.1 # Increased to unbreak canonicalize_version error (b/377774673)
 pytest-xdist==3.3.1 # Pinned to unbreak unit tests
 ray==2.5.0 # Pinned until 2.9.3 is verified for Ray tests
 ipython==8.22.2 # Pinned to unbreak TypeAliasType import error
-google-vizier==0.1.21
+google-vizier==0.1.21
+google-adk==0.0.2

testing/constraints-3.12.txt

@@ -11,3 +11,4 @@ packaging==24.1 # Increased to unbreak canonicalize_version error (b/377774673)
 pytest-xdist==3.3.1 # Pinned to unbreak unit tests
 ray==2.5.0 # Pinned until 2.9.3 is verified for Ray tests
 ipython==8.22.2 # Pinned to unbreak TypeAliasType import error
+google-adk==0.0.2

testing/constraints-3.8.txt

@@ -13,3 +13,4 @@ grpcio-testing==1.34.0
 pytest-xdist==3.3.1 # Pinned to unbreak unit tests
 ray==2.4.0 # Pinned until 2.9.3 is verified for Ray tests
 google-vizier==0.1.21
+google-adk==0.0.2

testing/constraints-3.9.txt

@@ -11,4 +11,5 @@ packaging==24.1 # Increased to unbreak canonicalize_version error (b/377774673)
 grpcio-testing==1.34.0
 pytest-xdist==3.3.1 # Pinned to unbreak unit tests
 ray==2.4.0 # Pinned until 2.9.3 is verified for Ray tests
-google-vizier==0.1.21
+google-vizier==0.1.21
+google-adk==0.0.2

tests/unit/architecture/test_vertexai_import.py

@@ -84,7 +84,7 @@ def test_vertexai_import():
     assert sorted(new_modules_after_vertexai) == [vertexai_module_name]
 
     assert vertexai_import_timedelta.total_seconds() < 0.005
-    assert aip_import_timedelta.total_seconds() < 23
+    assert aip_import_timedelta.total_seconds() < 40
 
     # Testing that external modules are not loaded.
     new_modules = modules_after_vertexai - modules_before_aip

tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py

@@ -0,0 +1,301 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import importlib
+import json
+from unittest import mock
+
+from google import auth
+import vertexai
+from google.cloud.aiplatform import initializer
+from vertexai.preview import reasoning_engines
+from vertexai.agent_engines import _utils
+import pytest
+
+
+try:
+    from google.adk.agents import llm_agent
+
+    Agent = llm_agent.Agent
+except ImportError:
+
+    class Agent:
+        def __init__(self, name: str, model: str):
+            self.name = name
+            self.model = model
+
+
+_TEST_LOCATION = "us-central1"
+_TEST_PROJECT = "test-project"
+_TEST_MODEL = "gemini-1.0-pro"
+
+
+@pytest.fixture(scope="module")
+def google_auth_mock():
+    with mock.patch.object(auth, "default") as google_auth_mock:
+        credentials_mock = mock.Mock()
+        credentials_mock.with_quota_project.return_value = None
+        google_auth_mock.return_value = (
+            credentials_mock,
+            _TEST_PROJECT,
+        )
+        yield google_auth_mock
+
+
+@pytest.fixture
+def vertexai_init_mock():
+    with mock.patch.object(vertexai, "init") as vertexai_init_mock:
+        yield vertexai_init_mock
+
+
+@pytest.fixture
+def cloud_trace_exporter_mock():
+    with mock.patch.object(
+        _utils,
+        "_import_cloud_trace_exporter_or_warn",
+    ) as cloud_trace_exporter_mock:
+        yield cloud_trace_exporter_mock
+
+
+@pytest.fixture
+def tracer_provider_mock():
+    with mock.patch("opentelemetry.sdk.trace.TracerProvider") as tracer_provider_mock:
+        yield tracer_provider_mock
+
+
+@pytest.fixture
+def simple_span_processor_mock():
+    with mock.patch(
+        "opentelemetry.sdk.trace.export.SimpleSpanProcessor"
+    ) as simple_span_processor_mock:
+        yield simple_span_processor_mock
+
+
+class _MockRunner:
+    def run(self, *args, **kwargs):
+        from google.adk.events import event
+
+        yield event.Event(
+            **{
+                "author": "currency_exchange_agent",
+                "content": {
+                    "parts": [
+                        {
+                            "function_call": {
+                                "args": {
+                                    "currency_date": "2025-04-03",
+                                    "currency_from": "USD",
+                                    "currency_to": "SEK",
+                                },
+                                "id": "af-c5a57692-9177-4091-a3df-098f834ee849",
+                                "name": "get_exchange_rate",
+                            }
+                        }
+                    ],
+                    "role": "model",
+                },
+                "id": "9aaItGK9",
+                "invocation_id": "e-6543c213-6417-484b-9551-b67915d1d5f7",
+            }
+        )
+
+
+@pytest.mark.usefixtures("google_auth_mock")
+class TestAdkApp:
+    def setup_method(self):
+        importlib.reload(initializer)
+        importlib.reload(vertexai)
+        vertexai.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
+
+    def teardown_method(self):
+        initializer.global_pool.shutdown(wait=True)
+
+    def test_initialization(self):
+        app = reasoning_engines.AdkApp(
+            agent=Agent(name="test_agent", model=_TEST_MODEL),
+        )
+        assert app._tmpl_attrs.get("project") == _TEST_PROJECT
+        assert app._tmpl_attrs.get("location") == _TEST_LOCATION
+        assert app._tmpl_attrs.get("runner") is None
+
+    def test_set_up(self):
+        app = reasoning_engines.AdkApp(
+            agent=Agent(name="test_agent", model=_TEST_MODEL),
+        )
+        assert app._tmpl_attrs.get("runner") is None
+        app.set_up()
+        assert app._tmpl_attrs.get("runner") is not None
+
+    def test_clone(self):
+        app = reasoning_engines.AdkApp(
+            agent=Agent(name="test_agent", model=_TEST_MODEL),
+        )
+        app.set_up()
+        assert app._tmpl_attrs.get("runner") is not None
+        app_clone = app.clone()
+        assert app._tmpl_attrs.get("runner") is not None
+        assert app_clone._tmpl_attrs.get("runner") is None
+        app_clone.set_up()
+        assert app_clone._tmpl_attrs.get("runner") is not None
+
+    def test_register_operations(self):
+        app = reasoning_engines.AdkApp(
+            agent=Agent(name="test_agent", model=_TEST_MODEL),
+        )
+        for operations in app.register_operations().values():
+            for operation in operations:
+                assert operation in dir(app)
+
+    def test_stream_query(self):
+        app = reasoning_engines.AdkApp(
+            agent=Agent(name="test_agent", model=_TEST_MODEL)
+        )
+        assert app._tmpl_attrs.get("runner") is None
+        app.set_up()
+        app._tmpl_attrs["runner"] = _MockRunner()
+        events = list(
+            app.stream_query(
+                user_id="test_user_id",
+                message="test message",
+            )
+        )
+        assert len(events) == 1
+
+    def test_streaming_agent_run_with_events(self):
+        app = reasoning_engines.AdkApp(
+            agent=Agent(name="test_agent", model=_TEST_MODEL)
+        )
+        app.set_up()
+        app._tmpl_attrs["in_memory_runner"] = _MockRunner()
+        request_json = json.dumps(
+            {
+                "artifacts": [
+                    {
+                        "file_name": "test_file_name",
+                        "versions": [{"version": "v1", "data": "v1data"}],
+                    }
+                ],
+                "authorizations": {
+                    "test_user_id1": {"access_token": "test_access_token"},
+                    "test_user_id2": {"accessToken": "test-access-token"},
+                },
+                "user_id": "test_user_id",
+                "message": {
+                    "parts": [{"text": "What is the exchange rate from USD to SEK?"}],
+                    "role": "user",
+                },
+            }
+        )
+        events = list(app.streaming_agent_run_with_events(request_json=request_json))
+        assert len(events) == 1
+
+    def test_create_session(self):
+        app = reasoning_engines.AdkApp(
+            agent=Agent(name="test_agent", model=_TEST_MODEL)
+        )
+        session1 = app.create_session(user_id="test_user_id")
+        assert session1.user_id == "test_user_id"
+        session2 = app.create_session(
+            user_id="test_user_id", session_id="test_session_id"
+        )
+        assert session2.user_id == "test_user_id"
+        assert session2.id == "test_session_id"
+
+    def test_get_session(self):
+        app = reasoning_engines.AdkApp(
+            agent=Agent(name="test_agent", model=_TEST_MODEL)
+        )
+        session1 = app.create_session(user_id="test_user_id")
+        session2 = app.get_session(
+            user_id="test_user_id",
+            session_id=session1.id,
+        )
+        assert session2.user_id == "test_user_id"
+        assert session1.id == session2.id
+
+    def test_list_sessions(self):
+        app = reasoning_engines.AdkApp(
+            agent=Agent(name="test_agent", model=_TEST_MODEL)
+        )
+        response0 = app.list_sessions(user_id="test_user_id")
+        assert not response0.sessions
+        session = app.create_session(user_id="test_user_id")
+        response1 = app.list_sessions(user_id="test_user_id")
+        assert len(response1.sessions) == 1
+        assert response1.sessions[0].id == session.id
+        session2 = app.create_session(user_id="test_user_id")
+        response2 = app.list_sessions(user_id="test_user_id")
+        assert len(response2.sessions) == 2
+        assert response2.sessions[0].id == session.id
+        assert response2.sessions[1].id == session2.id
+
+    def test_delete_session(self):
+        app = reasoning_engines.AdkApp(
+            agent=Agent(name="test_agent", model=_TEST_MODEL)
+        )
+        response = app.delete_session(user_id="test_user_id", session_id="")
+        assert not response
+        session = app.create_session(user_id="test_user_id")
+        response1 = app.list_sessions(user_id="test_user_id")
+        assert len(response1.sessions) == 1
+        app.delete_session(user_id="test_user_id", session_id=session.id)
+        response0 = app.list_sessions(user_id="test_user_id")
+        assert not response0.sessions
+
+    @pytest.mark.usefixtures("caplog")
+    def test_enable_tracing(
+        self,
+        caplog,
+        cloud_trace_exporter_mock,
+        tracer_provider_mock,
+        simple_span_processor_mock,
+    ):
+        app = reasoning_engines.AdkApp(
+            agent=Agent(name="test_agent", model=_TEST_MODEL),
+            enable_tracing=True,
+        )
+        assert app._tmpl_attrs.get("instrumentor") is None
+        # TODO(b/384730642): Re-enable this test once the parent issue is fixed.
+        # agent.set_up()
+        # assert agent._tmpl_attrs.get("instrumentor") is not None
+        # assert (
+        #     "enable_tracing=True but proceeding with tracing disabled"
+        #     not in caplog.text
+        # )
+
+    @pytest.mark.usefixtures("caplog")
+    def test_enable_tracing_warning(self, caplog):
+        app = reasoning_engines.AdkApp(
+            agent=Agent(name="test_agent", model=_TEST_MODEL),
+            enable_tracing=True,
+        )
+        assert app._tmpl_attrs.get("instrumentor") is None
+        # TODO(b/384730642): Re-enable this test once the parent issue is fixed.
+        # app.set_up()
+        # assert "enable_tracing=True but proceeding with tracing disabled" in caplog.text
+
+
+class TestAdkAppErrors:
+    def test_raise_get_session_not_found_error(self):
+        with pytest.raises(
+            RuntimeError,
+            match=r"Session not found. Please create it using .create_session()",
+        ):
+            app = reasoning_engines.AdkApp(
+                agent=Agent(name="test_agent", model=_TEST_MODEL),
+            )
+            app.get_session(
+                user_id="non_existent_user",
+                session_id="test_session_id",
+            )