Skip to content

Commit faef621

Browse files
vertex-sdk-botcopybara-github
authored andcommitted
chore: pre-emptive v2-compatible test/code fixes for ADK migration
PiperOrigin-RevId: 910099643
1 parent a2654bb commit faef621

2 files changed

Lines changed: 46 additions & 50 deletions

File tree

tests/unit/vertex_adk/test_agent_engine_templates_adk.py

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,7 @@ def __init__(self, name: str, model: str):
6363
_TEST_AGENT = Agent(name=_TEST_AGENT_NAME, model=_TEST_MODEL)
6464
_TEST_SESSION = {
6565
"id": "ca18c25a-644b-4e13-9b24-78c150ec3eb9",
66-
"app_name": "default-app-name",
66+
"app_name": "default_app_name",
6767
"user_id": _TEST_USER_ID,
6868
"events": [
6969
{
@@ -325,8 +325,6 @@ def test_adk_version(self):
325325
agent_engines.AdkApp(agent=_TEST_AGENT)
326326

327327
def setup_method(self):
328-
importlib.reload(initializer)
329-
importlib.reload(vertexai)
330328
vertexai.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
331329

332330
def teardown_method(self):
@@ -1067,8 +1065,6 @@ def update_agent_engine_mock():
10671065
@pytest.mark.usefixtures("google_auth_mock")
10681066
class TestAgentEngines:
10691067
def setup_method(self):
1070-
importlib.reload(initializer)
1071-
importlib.reload(aiplatform)
10721068
aiplatform.init(
10731069
project=_TEST_PROJECT,
10741070
location=_TEST_LOCATION,

vertexai/agent_engines/templates/adk.py

Lines changed: 45 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -32,68 +32,68 @@
3232
import warnings
3333

3434
if TYPE_CHECKING:
35-
try:
36-
from google.adk.events.event import Event
35+
try:
36+
from google.adk.events.event import Event
3737

38-
Event = Event
39-
except (ImportError, AttributeError):
40-
Event = Any
38+
Event = Event
39+
except (ImportError, AttributeError):
40+
Event = Any
4141

42-
try:
43-
from google.adk.apps import App
42+
try:
43+
from google.adk.apps import App
4444

45-
App = App
46-
except (ImportError, AttributeError):
47-
App = Any
45+
App = App
46+
except (ImportError, AttributeError):
47+
App = Any
4848

49-
try:
50-
from google.adk.agents import BaseAgent
49+
try:
50+
from google.adk.agents import BaseAgent
5151

52-
BaseAgent = BaseAgent
53-
except (ImportError, AttributeError):
54-
BaseAgent = Any
52+
BaseAgent = BaseAgent
53+
except (ImportError, AttributeError):
54+
BaseAgent = Any
5555

56-
try:
57-
from google.adk.plugins.base_plugin import BasePlugin
56+
try:
57+
from google.adk.plugins.base_plugin import BasePlugin
5858

59-
BasePlugin = BasePlugin
60-
except (ImportError, AttributeError):
61-
BasePlugin = Any
59+
BasePlugin = BasePlugin
60+
except (ImportError, AttributeError):
61+
BasePlugin = Any
6262

63-
try:
64-
from google.adk.sessions import BaseSessionService
63+
try:
64+
from google.adk.sessions import BaseSessionService
6565

66-
BaseSessionService = BaseSessionService
67-
except (ImportError, AttributeError):
68-
BaseSessionService = Any
66+
BaseSessionService = BaseSessionService
67+
except (ImportError, AttributeError):
68+
BaseSessionService = Any
6969

70-
try:
71-
from google.adk.artifacts import BaseArtifactService
70+
try:
71+
from google.adk.artifacts import BaseArtifactService
7272

73-
BaseArtifactService = BaseArtifactService
74-
except (ImportError, AttributeError):
75-
BaseArtifactService = Any
73+
BaseArtifactService = BaseArtifactService
74+
except (ImportError, AttributeError):
75+
BaseArtifactService = Any
7676

77-
try:
78-
from google.adk.memory import BaseMemoryService
77+
try:
78+
from google.adk.memory import BaseMemoryService
7979

80-
BaseMemoryService = BaseMemoryService
81-
except (ImportError, AttributeError):
82-
BaseMemoryService = Any
80+
BaseMemoryService = BaseMemoryService
81+
except (ImportError, AttributeError):
82+
BaseMemoryService = Any
8383

84-
try:
85-
from opentelemetry.sdk import trace
84+
try:
85+
from opentelemetry.sdk import trace
8686

87-
TracerProvider = trace.TracerProvider
88-
SpanProcessor = trace.SpanProcessor
89-
SynchronousMultiSpanProcessor = trace.SynchronousMultiSpanProcessor
90-
except (ImportError, AttributeError):
91-
TracerProvider = Any
92-
SpanProcessor = Any
93-
SynchronousMultiSpanProcessor = Any
87+
TracerProvider = trace.TracerProvider
88+
SpanProcessor = trace.SpanProcessor
89+
SynchronousMultiSpanProcessor = trace.SynchronousMultiSpanProcessor
90+
except (ImportError, AttributeError):
91+
TracerProvider = Any
92+
SpanProcessor = Any
93+
SynchronousMultiSpanProcessor = Any
9494

9595

96-
_DEFAULT_APP_NAME = "default-app-name"
96+
_DEFAULT_APP_NAME = "default_app_name"
9797
_DEFAULT_USER_ID = "default-user-id"
9898
_TELEMETRY_API_DISABLED_WARNING = (
9999
"Tracing integration for Agent Engine has migrated to a new API.\n"

0 commit comments

Comments
 (0)