Skip to content

Implement lazy resolution for project #4465

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
jayy-77 wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Implement lazy resolution for project #4465

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 4 additions & 6 deletions src/google/adk/cli/service_registry.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -357,15 +357,12 @@ def _parse_agent_engine_kwargs(
"Agent engine resource name or resource id cannot be empty."
)

# If uri_part is just an ID, load project/location from env
# If uri_part is just an ID, defer project/location loading to runtime
if "/" not in uri_part:
project, location = _load_gcp_config(
agents_dir, "short-form agent engine IDs"
)
# Return with agents_dir for lazy resolution at runtime
return {
"project": project,
"location": location,
"agent_engine_id": uri_part,
"agents_dir": agents_dir,
}

# If uri_part is a full resource name, parse it
Expand All @@ -385,6 +382,7 @@ def _parse_agent_engine_kwargs(
"project": parts[1],
"location": parts[3],
"agent_engine_id": parts[5],
"agents_dir": agents_dir,
}


Expand Down
69 changes: 63 additions & 6 deletions src/google/adk/memory/vertex_ai_memory_bank_service.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -72,12 +72,15 @@ def __init__(
agent_engine_id: Optional[str] = None,
*,
express_mode_api_key: Optional[str] = None,
agents_dir: Optional[str] = None,
):
"""Initializes a VertexAiMemoryBankService.

Args:
project: The project ID of the Memory Bank to use.
location: The location of the Memory Bank to use.
project: The project ID of the Memory Bank to use. If not provided, will
be resolved lazily at runtime from environment variables or .env files.
location: The location of the Memory Bank to use. If not provided, will
be resolved lazily at runtime from environment variables or .env files.
agent_engine_id: The ID of the agent engine to use for the Memory Bank,
e.g. '456' in
'projects/my-project/locations/us-central1/reasoningEngines/456'. To
Expand All @@ -88,13 +91,15 @@ def __init__(
be used. It will only be used if GOOGLE_GENAI_USE_VERTEXAI is true. Do
not use Google AI Studio API key for this field. For more details, visit
https://cloud.google.com/vertex-ai/generative-ai/docs/start/express-mode/overview
agents_dir: The directory containing agent configurations and .env files.
Used for lazy resolution of project/location when not explicitly provided.
"""
self._project = project
self._location = location
self._agent_engine_id = agent_engine_id
self._express_mode_api_key = get_express_mode_api_key(
project, location, express_mode_api_key
)
self._agents_dir = agents_dir
self._config_resolved = False
self._express_mode_api_key = express_mode_api_key

if agent_engine_id and '/' in agent_engine_id:
logger.warning(
Expand Down Expand Up @@ -168,6 +173,9 @@ async def _add_events_to_memory_from_events(

@override
async def search_memory(self, *, app_name: str, user_id: str, query: str):
# Lazily resolve project/location on first use
self._resolve_config()

if not self._agent_engine_id:
raise ValueError('Agent Engine ID is required for Memory Bank.')

Expand Down Expand Up @@ -203,7 +211,56 @@ async def search_memory(self, *, app_name: str, user_id: str, query: str):
)
return SearchMemoryResponse(memories=memory_events)

def _get_api_client(self) -> vertexai.AsyncClient:
def _resolve_config(self) -> None:
"""Lazily resolves project and location if not provided at initialization.

This method is called on first use to resolve GCP configuration from:
1. Explicit environment variables (highest priority)
2. agents_dir root .env file
3. Parent directory .env files (walking upward)

Raises:
ValueError: If project or location cannot be resolved.
"""
if self._config_resolved:
return

import os

# If both are already set (either at init or via environment), we're done
if self._project and self._location:
self._config_resolved = True
self._express_mode_api_key = get_express_mode_api_key(
self._project, self._location, self._express_mode_api_key
)
return

# Try to load from environment and .env files
if self._agents_dir:
# Load from agents_dir root
from ..cli.utils import envs
envs.load_dotenv_for_agent("", self._agents_dir)

# Resolve from environment after loading .env
self._project = self._project or os.environ.get("GOOGLE_CLOUD_PROJECT")
self._location = self._location or os.environ.get("GOOGLE_CLOUD_LOCATION")

if not self._project or not self._location:
error_msg = (
"GOOGLE_CLOUD_PROJECT and GOOGLE_CLOUD_LOCATION must be set. "
"You can set them via:\n"
" 1. Environment variables: GOOGLE_CLOUD_PROJECT and GOOGLE_CLOUD_LOCATION\n"
" 2. .env file in your agents_dir\n"
" 3. Use full resource name: agentengine://projects/{project}/locations/{location}/reasoningEngines/{id}"
)
raise ValueError(error_msg)

self._config_resolved = True
self._express_mode_api_key = get_express_mode_api_key(
self._project, self._location, self._express_mode_api_key
)
Comment on lines +214 to +261
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Feb 12, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The _resolve_config method has some repeated logic for setting self._config_resolved and calling get_express_mode_api_key. This can be refactored to be more concise and follow the DRY (Don't Repeat Yourself) principle.

  def _resolve_config(self) -> None:
    """Lazily resolves project and location if not provided at initialization.
    
    This method is called on first use to resolve GCP configuration from:
    1. Explicit environment variables (highest priority)
    2. agents_dir root .env file
    3. Parent directory .env files (walking upward)
    
    Raises:
      ValueError: If project or location cannot be resolved.
    """
    if self._config_resolved:
      return
    
    import os
    
    if not (self._project and self._location):
      # Try to load from environment and .env files
      if self._agents_dir:
        # Load from agents_dir root
        from ..cli.utils import envs
        envs.load_dotenv_for_agent("", self._agents_dir)
      
      # Resolve from environment after loading .env
      self._project = self._project or os.environ.get("GOOGLE_CLOUD_PROJECT")
      self._location = self._location or os.environ.get("GOOGLE_CLOUD_LOCATION")
    
    if not self._project or not self._location:
      error_msg = (
          "GOOGLE_CLOUD_PROJECT and GOOGLE_CLOUD_LOCATION must be set. "
          "You can set them via:\n"
          "  1. Environment variables: GOOGLE_CLOUD_PROJECT and GOOGLE_CLOUD_LOCATION\n"
          "  2. .env file in your agents_dir\n"
          "  3. Use full resource name: agentengine://projects/{project}/locations/{location}/reasoningEngines/{id}"
      )
      raise ValueError(error_msg)
    
    self._config_resolved = True
    self._express_mode_api_key = get_express_mode_api_key(
        self._project, self._location, self._express_mode_api_key
    )


def _get_api_client(self):
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Feb 12, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The return type hint for _get_api_client was removed in this change. Restoring it will improve code clarity and help static analysis tools. The method returns an AsyncClient.

  def _get_api_client(self) -> vertexai.AsyncClient:

"""Instantiates an API client for the given project and location.

It needs to be instantiated inside each request so that the event loop
Expand Down
87 changes: 82 additions & 5 deletions src/google/adk/sessions/vertex_ai_session_service.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,26 +55,31 @@ def __init__(
agent_engine_id: Optional[str] = None,
*,
express_mode_api_key: Optional[str] = None,
agents_dir: Optional[str] = None,
):
"""Initializes the VertexAiSessionService.

Args:
project: The project id of the project to use.
location: The location of the project to use.
project: The project id of the project to use. If not provided, will be
resolved lazily at runtime from environment variables or .env files.
location: The location of the project to use. If not provided, will be
resolved lazily at runtime from environment variables or .env files.
agent_engine_id: The resource ID of the agent engine to use.
express_mode_api_key: The API key to use for Express Mode. If not
provided, the API key from the GOOGLE_API_KEY environment variable will
be used. It will only be used if GOOGLE_GENAI_USE_VERTEXAI is true.
Do not use Google AI Studio API key for this field. For more details,
visit
https://cloud.google.com/vertex-ai/generative-ai/docs/start/express-mode/overview
agents_dir: The directory containing agent configurations and .env files.
Used for lazy resolution of project/location when not explicitly provided.
"""
self._project = project
self._location = location
self._agent_engine_id = agent_engine_id
self._express_mode_api_key = get_express_mode_api_key(
project, location, express_mode_api_key
)
self._agents_dir = agents_dir
self._config_resolved = False
self._express_mode_api_key = express_mode_api_key

@override
async def create_session(
Expand All @@ -100,6 +105,8 @@ async def create_session(
Returns:
The created session.
"""
# Lazily resolve project/location on first use
self._resolve_config(app_name)

if session_id:
raise ValueError(
Expand Down Expand Up @@ -139,6 +146,9 @@ async def get_session(
session_id: str,
config: Optional[GetSessionConfig] = None,
) -> Optional[Session]:
# Lazily resolve project/location on first use
self._resolve_config(app_name)

reasoning_engine_id = self._get_reasoning_engine_id(app_name)
session_resource_name = (
f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}'
Expand Down Expand Up @@ -203,6 +213,9 @@ async def get_session(
async def list_sessions(
self, *, app_name: str, user_id: Optional[str] = None
) -> ListSessionsResponse:
# Lazily resolve project/location on first use
self._resolve_config(app_name)

reasoning_engine_id = self._get_reasoning_engine_id(app_name)

async with self._get_api_client() as api_client:
Expand Down Expand Up @@ -231,6 +244,9 @@ async def list_sessions(
async def delete_session(
self, *, app_name: str, user_id: str, session_id: str
) -> None:
# Lazily resolve project/location on first use
self._resolve_config(app_name)

reasoning_engine_id = self._get_reasoning_engine_id(app_name)

async with self._get_api_client() as api_client:
Expand All @@ -249,6 +265,9 @@ async def append_event(self, session: Session, event: Event) -> Event:
# Update the in-memory session.
await super().append_event(session=session, event=event)

# Lazily resolve project/location on first use
self._resolve_config(session.app_name)

reasoning_engine_id = self._get_reasoning_engine_id(session.app_name)

config = {}
Expand Down Expand Up @@ -323,6 +342,64 @@ def _get_reasoning_engine_id(self, app_name: str):

return match.groups()[-1]

def _resolve_config(self, app_name: Optional[str] = None) -> None:
"""Lazily resolves project and location if not provided at initialization.

This method is called on first use to resolve GCP configuration from:
1. Explicit environment variables (highest priority)
2. Agent-specific .env file (if app_name and agents_dir provided)
3. agents_dir root .env file
4. Parent directory .env files (walking upward)

Args:
app_name: Optional app name to load agent-specific .env files.

Raises:
ValueError: If project or location cannot be resolved.
"""
if self._config_resolved:
return

import os

# If both are already set (either at init or via environment), we're done
if self._project and self._location:
self._config_resolved = True
self._express_mode_api_key = get_express_mode_api_key(
self._project, self._location, self._express_mode_api_key
)
return

# Try to load from environment and .env files
if self._agents_dir and app_name:
# Load agent-specific .env
from ..cli.utils import envs
envs.load_dotenv_for_agent(app_name, self._agents_dir)
elif self._agents_dir:
# Load from agents_dir root
from ..cli.utils import envs
envs.load_dotenv_for_agent("", self._agents_dir)

# Resolve from environment after loading .env
self._project = self._project or os.environ.get("GOOGLE_CLOUD_PROJECT")
self._location = self._location or os.environ.get("GOOGLE_CLOUD_LOCATION")

if not self._project or not self._location:
error_msg = (
"GOOGLE_CLOUD_PROJECT and GOOGLE_CLOUD_LOCATION must be set. "
"You can set them via:\n"
" 1. Environment variables: GOOGLE_CLOUD_PROJECT and GOOGLE_CLOUD_LOCATION\n"
" 2. .env file in your agent directory\n"
" 3. .env file in your agents_dir\n"
" 4. Use full resource name: agentengine://projects/{project}/locations/{location}/reasoningEngines/{id}"
)
raise ValueError(error_msg)

self._config_resolved = True
self._express_mode_api_key = get_express_mode_api_key(
self._project, self._location, self._express_mode_api_key
)
Comment on lines +345 to +401
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Feb 12, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

This _resolve_config method has some repeated logic, similar to the one in VertexAiMemoryBankService. Additionally, the .env loading part can be cleaned up to avoid a repeated import. Refactoring this will improve readability and maintainability.

  def _resolve_config(self, app_name: Optional[str] = None) -> None:
    """Lazily resolves project and location if not provided at initialization.
    
    This method is called on first use to resolve GCP configuration from:
    1. Explicit environment variables (highest priority)
    2. Agent-specific .env file (if app_name and agents_dir provided)
    3. agents_dir root .env file
    4. Parent directory .env files (walking upward)
    
    Args:
      app_name: Optional app name to load agent-specific .env files.
    
    Raises:
      ValueError: If project or location cannot be resolved.
    """
    if self._config_resolved:
      return
    
    import os
    
    if not (self._project and self._location):
      # Try to load from environment and .env files
      if self._agents_dir:
        from ..cli.utils import envs
        if app_name:
          # Load agent-specific .env
          envs.load_dotenv_for_agent(app_name, self._agents_dir)
        else:
          # Load from agents_dir root
          envs.load_dotenv_for_agent("", self._agents_dir)
      
      # Resolve from environment after loading .env
      self._project = self._project or os.environ.get("GOOGLE_CLOUD_PROJECT")
      self._location = self._location or os.environ.get("GOOGLE_CLOUD_LOCATION")
    
    if not self._project or not self._location:
      error_msg = (
          "GOOGLE_CLOUD_PROJECT and GOOGLE_CLOUD_LOCATION must be set. "
          "You can set them via:\n"
          "  1. Environment variables: GOOGLE_CLOUD_PROJECT and GOOGLE_CLOUD_LOCATION\n"
          "  2. .env file in your agent directory\n"
          "  3. .env file in your agents_dir\n"
          "  4. Use full resource name: agentengine://projects/{project}/locations/{location}/reasoningEngines/{id}"
      )
      raise ValueError(error_msg)
    
    self._config_resolved = True
    self._express_mode_api_key = get_express_mode_api_key(
        self._project, self._location, self._express_mode_api_key
    )


def _api_client_http_options_override(
self,
) -> Optional[Union[types.HttpOptions, types.HttpOptionsDict]]:
Expand Down