Skip to content

Commit ec97a58

Browse files
vertex-sdk-botcopybara-github
authored andcommitted
feat: Add support for API key in AdkApp
PiperOrigin-RevId: 825638989
1 parent 33b3d2b commit ec97a58

File tree

1 file changed

+56
-24
lines changed
  • vertexai/agent_engines/templates

1 file changed

+56
-24
lines changed

vertexai/agent_engines/templates/adk.py

Lines changed: 56 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -505,13 +505,13 @@ def __init__(
505505
"""
506506
from google.cloud.aiplatform import initializer
507507

508-
adk_version = get_adk_version()
509-
if not is_version_sufficient("1.5.0"):
510-
msg = (
511-
f"Unsupported google-adk version: {adk_version}, please use "
512-
"google-adk>=1.5.0 for AdkApp deployment on Agent Engine."
513-
)
514-
raise ValueError(msg)
508+
# adk_version = get_adk_version()
509+
# if not is_version_sufficient("1.5.0"):
510+
# msg = (
511+
# f"Unsupported google-adk version: {adk_version}, please use "
512+
# "google-adk>=1.5.0 for AdkApp deployment on Agent Engine."
513+
# )
514+
# raise ValueError(msg)
515515

516516
if not agent and not app:
517517
raise ValueError("One of `agent` or `app` must be provided.")
@@ -540,6 +540,7 @@ def __init__(
540540
"artifact_service_builder": artifact_service_builder,
541541
"memory_service_builder": memory_service_builder,
542542
"instrumentor_builder": instrumentor_builder,
543+
"express_mode_api_key": initializer.global_config.api_key,
543544
}
544545

545546
async def _init_session(
@@ -683,9 +684,16 @@ def set_up(self):
683684

684685
os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "1"
685686
project = self._tmpl_attrs.get("project")
686-
os.environ["GOOGLE_CLOUD_PROJECT"] = project
687+
if project:
688+
os.environ["GOOGLE_CLOUD_PROJECT"] = project
687689
location = self._tmpl_attrs.get("location")
688-
os.environ["GOOGLE_CLOUD_LOCATION"] = location
690+
if location:
691+
os.environ["GOOGLE_CLOUD_LOCATION"] = location
692+
express_mode_api_key = self._tmpl_attrs.get("express_mode_api_key")
693+
if express_mode_api_key and not project and not location:
694+
os.environ["GOOGLE_API_KEY"] = express_mode_api_key
695+
if not project:
696+
os.environ.pop("GOOGLE_CLOUD_LOCATION", None)
689697

690698
# Disable content capture in custom ADK spans unless user enabled
691699
# tracing explicitly with the old flag
@@ -769,21 +777,37 @@ def tracing_enabled() -> bool:
769777
VertexAiSessionService,
770778
)
771779

772-
self._tmpl_attrs["session_service"] = VertexAiSessionService(
773-
project=project,
774-
location=location,
775-
agent_engine_id=os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID"),
776-
)
780+
if is_version_sufficient("1.18.0"):
781+
self._tmpl_attrs["session_service"] = VertexAiSessionService(
782+
project=project,
783+
location=location,
784+
agent_engine_id=os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID"),
785+
express_mode_api_key=express_mode_api_key,
786+
)
787+
else:
788+
self._tmpl_attrs["session_service"] = VertexAiSessionService(
789+
project=project,
790+
location=location,
791+
agent_engine_id=os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID"),
792+
)
777793
except (ImportError, AttributeError):
778794
from google.adk.sessions.vertex_ai_session_service_g3 import (
779795
VertexAiSessionService,
780796
)
781797

782-
self._tmpl_attrs["session_service"] = VertexAiSessionService(
783-
project=project,
784-
location=location,
785-
agent_engine_id=os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID"),
786-
)
798+
if is_version_sufficient("1.18.0"):
799+
self._tmpl_attrs["session_service"] = VertexAiSessionService(
800+
project=project,
801+
location=location,
802+
agent_engine_id=os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID"),
803+
express_mode_api_key=express_mode_api_key,
804+
)
805+
else:
806+
self._tmpl_attrs["session_service"] = VertexAiSessionService(
807+
project=project,
808+
location=location,
809+
agent_engine_id=os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID"),
810+
)
787811

788812
else:
789813
self._tmpl_attrs["session_service"] = InMemorySessionService()
@@ -799,11 +823,19 @@ def tracing_enabled() -> bool:
799823
VertexAiMemoryBankService,
800824
)
801825

802-
self._tmpl_attrs["memory_service"] = VertexAiMemoryBankService(
803-
project=project,
804-
location=location,
805-
agent_engine_id=os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID"),
806-
)
826+
if is_version_sufficient("1.18.0"):
827+
self._tmpl_attrs["memory_service"] = VertexAiMemoryBankService(
828+
project=project,
829+
location=location,
830+
agent_engine_id=os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID"),
831+
express_mode_api_key=express_mode_api_key,
832+
)
833+
else:
834+
self._tmpl_attrs["memory_service"] = VertexAiMemoryBankService(
835+
project=project,
836+
location=location,
837+
agent_engine_id=os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID"),
838+
)
807839
except (ImportError, AttributeError):
808840
# TODO(ysian): Handle this via _g3 import for google3.
809841
pass

0 commit comments

Comments
 (0)