fix: Update definition of async_add_session_to_memory to not depend on Session

This avoids issues with the generation of the schema when deploying PiperOrigin-RevId: 808304051

Author: yeesian

vertexai/_genai/_agent_engines_utils.py

@@ -91,21 +91,6 @@
  _SpecifierSet: type[Any] = Any
 
 
-try:
- # For the registration of the AdkApp.async_add_session_to_memory method to
- # avoid errors of the form:
- # failed to generate schema for async_add_session_to_memory:
- # `async_add_session_to_memory` is not fully defined; you should define
- # `Session`, then call `async_add_session_to_memory.model_rebuild()`
- from google.adk.sessions.session import Session
-
- Session = Session
-except (ImportError, AttributeError):
- # For agent engine templates that do not depend on ADK, we do not need to
- # worry about registering the AdkApp.async_add_session_to_memory method.
- Session = Any
-
-
 try:
  from a2a.types import (
  AgentCard,

vertexai/agent_engines/templates/adk.py

@@ -53,13 +53,6 @@
  except (ImportError, AttributeError):
  BaseSessionService = Any
 
- try:
- from google.adk.sessions.session import Session
-
- Session = Session
- except (ImportError, AttributeError):
- Session = Any
-
  try:
  from google.adk.artifacts import BaseArtifactService
 
@@ -842,16 +835,14 @@ async def async_delete_session(
  **kwargs,
  )
 
- async def async_add_session_to_memory(
- self,
- *,
- session: Union["Session", Dict[str, Any]],
- ):
+ async def async_add_session_to_memory(self, *, session: Dict[str, Any]):
  """Generates memories.
 
  Args:
- session (Union[Session, Dict[str, Any]]):
- Required. The session to use for generating memories.
+ session (Dict[str, Any]):
+ Required. The session to use for generating memories. It should
+ be a dictionary representing an ADK Session object, e.g.
+ session.model_dump(mode="json").
  """
  from google.adk.sessions.session import Session
 

vertexai/preview/reasoning_engines/templates/adk.py

@@ -58,13 +58,6 @@
  except (ImportError, AttributeError):
  BaseSessionService = Any
 
- try:
- from google.adk.sessions.session import Session
-
- Session = Session
- except (ImportError, AttributeError):
- Session = Any
-
  try:
  from google.adk.artifacts import BaseArtifactService
 
@@ -1135,16 +1128,14 @@ def _asyncio_thread_main():
  if isinstance(outcome, RuntimeError):
  raise outcome from None
 
- async def async_add_session_to_memory(
- self,
- *,
- session: Union["Session", Dict[str, Any]],
- ):
+ async def async_add_session_to_memory(self, *, session: Dict[str, Any]):
  """Generates memories.
 
  Args:
- session (Union[Session, Dict[str, Any]]):
- Required. The session to use for generating memories.
+ session (Dict[str, Any]):
+ Required. The session to use for generating memories. It should
+ be a dictionary representing an ADK Session object, e.g.
+ session.model_dump(mode="json").
  """
  from google.adk.sessions.session import Session