Skip to content

Commit cfb44c4

Browse files
vertex-sdk-botcopybara-github
authored andcommitted
fix: Keep AdkApp subclass, plugins and credential_service_builder in clone()
PiperOrigin-RevId: 986375881
1 parent 6b29d01 commit cfb44c4

2 files changed

Lines changed: 30 additions & 1 deletion

File tree

‎tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -418,6 +418,31 @@ def test_clone(self):
418418
app_clone.set_up()
419419
assert app_clone._tmpl_attrs.get("runner") is not None
420420

421+
def test_clone_preserves_subclass(self):
422+
class CustomAdkApp(reasoning_engines.AdkApp):
423+
pass
424+
425+
app = CustomAdkApp(
426+
agent=Agent(name=_TEST_AGENT_NAME, model=_TEST_MODEL),
427+
)
428+
assert isinstance(app.clone(), CustomAdkApp)
429+
430+
def test_clone_preserves_plugins_and_credential_service_builder(self):
431+
def credential_service_builder():
432+
return None
433+
434+
app = reasoning_engines.AdkApp(
435+
agent=Agent(name=_TEST_AGENT_NAME, model=_TEST_MODEL),
436+
plugins=["test_plugin"],
437+
credential_service_builder=credential_service_builder,
438+
)
439+
app_clone = app.clone()
440+
assert app_clone._tmpl_attrs.get("plugins") == ["test_plugin"]
441+
assert (
442+
app_clone._tmpl_attrs.get("credential_service_builder")
443+
is credential_service_builder
444+
)
445+
421446
def test_register_operations(self):
422447
app = reasoning_engines.AdkApp(
423448
agent=Agent(name=_TEST_AGENT_NAME, model=_TEST_MODEL),

‎vertexai/preview/reasoning_engines/templates/adk.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -846,12 +846,16 @@ def clone(self):
846846
"""Returns a clone of the ADK application."""
847847
import copy
848848

849-
return AdkApp(
849+
return self.__class__(
850850
agent=copy.deepcopy(self._tmpl_attrs.get("agent")),
851+
plugins=copy.deepcopy(self._tmpl_attrs.get("plugins")),
851852
enable_tracing=self._tmpl_attrs.get("enable_tracing"),
852853
session_service_builder=self._tmpl_attrs.get("session_service_builder"),
853854
artifact_service_builder=self._tmpl_attrs.get("artifact_service_builder"),
854855
memory_service_builder=self._tmpl_attrs.get("memory_service_builder"),
856+
credential_service_builder=self._tmpl_attrs.get(
857+
"credential_service_builder"
858+
),
855859
env_vars=self._tmpl_attrs.get("env_vars"),
856860
)
857861

0 commit comments

Comments
 (0)