feat: Add memory related methods to AdkApp

PiperOrigin-RevId: 794137925

Author: yeesian

tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py

@@ -370,6 +370,105 @@ def test_delete_session(self):
  response0 = app.list_sessions(user_id="test_user_id")
  assert not response0.sessions
 
+ @pytest.mark.asyncio
+ async def test_async_add_session_to_memory(self):
+ app = reasoning_engines.AdkApp(
+ agent=Agent(name="test_agent", model=_TEST_MODEL)
+ )
+ assert app._tmpl_attrs.get("memory_service") is None
+ session = app.create_session(user_id="test_user_id")
+ list(
+ app.stream_query(
+ user_id="test_user_id",
+ session_id=session.id,
+ message="My cat's name is Garfield",
+ )
+ )
+ await app.async_add_session_to_memory(
+ session=app.get_session(
+ user_id="test_user_id",
+ session_id=session.id,
+ )
+ )
+ response = await app.async_search_memory(
+ user_id="test_user_id",
+ query="What is my cat's name",
+ )
+ assert len(response.memories) >= 1
+
+ @pytest.mark.asyncio
+ async def test_async_add_session_to_memory_dict(self):
+ app = reasoning_engines.AdkApp(
+ agent=Agent(name="test_agent", model=_TEST_MODEL)
+ )
+ await app.async_add_session_to_memory(
+ session={
+ "id": "ca18c25a-644b-4e13-9b24-78c150ec3eb9",
+ "app_name": "default-app-name",
+ "user_id": "test_user_id",
+ "events": [
+ {
+ "author": "user",
+ "content": {
+ "parts": [{"text": "My cat's name is Garfield"}],
+ "role": "user",
+ },
+ },
+ {
+ "author": "my_personal_agent",
+ "content": {
+ "parts": [{"text": "Okay, good to know!"}],
+ "role": "model",
+ },
+ },
+ ],
+ },
+ )
+ response = await app.async_search_memory(
+ user_id="test_user_id",
+ query="What is my cat's name",
+ )
+ assert len(response.memories) >= 1
+
+ @pytest.mark.asyncio
+ async def test_async_search_memory(self):
+ app = reasoning_engines.AdkApp(
+ agent=Agent(name="test_agent", model=_TEST_MODEL)
+ )
+ response = await app.async_search_memory(
+ user_id="test_user_id",
+ query="What is my cat's name",
+ )
+ assert not response.memories
+ await app.async_add_session_to_memory(
+ session={
+ "id": "ca18c25a-644b-4e13-9b24-78c150ec3eb9",
+ "app_name": "default-app-name",
+ "user_id": "test_user_id",
+ "events": [
+ {
+ "author": "user",
+ "content": {
+ "parts": [{"text": "My cat's name is Garfield"}],
+ "role": "user",
+ },
+ },
+ {
+ "author": "my_personal_agent",
+ "content": {
+ "parts": [{"text": "Okay, good to know!"}],
+ "role": "model",
+ },
+ },
+ ],
+ },
+ )
+ response = await app.async_search_memory(
+ user_id="test_user_id",
+ query="What is my cat's name",
+ )
+ assert len(response.memories) >= 1
+
  @pytest.mark.usefixtures("caplog")
  def test_enable_tracing(
  self,

vertexai/preview/reasoning_engines/templates/adk.py

@@ -51,6 +51,13 @@
  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
 
@@ -961,6 +968,53 @@ 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]],
+ ):
+ """Generates memories.
+
+ Args:
+ session (Union[Session, Dict[str, Any]]):
+ Required. The session to use for generating memories.
+ """
+ from google.adk.sessions.session import Session
+
+ if isinstance(session, Dict):
+ session = Session.model_validate(session)
+ elif not isinstance(session, Session):
+ raise TypeError("session must be a Session object.")
+ if not session.events:
+ # Get the latest version of the session in case it was updated.
+ session = await self.async_get_session(
+ user_id=session.user_id,
+ session_id=session.id,
+ )
+ if not self._tmpl_attrs.get("memory_service"):
+ self.set_up()
+ return await self._tmpl_attrs.get("memory_service").add_session_to_memory(
+ session=session,
+ )
+
+ async def async_search_memory(self, *, user_id: str, query: str):
+ """Searches memories for the given user.
+
+ Args:
+ user_id: The id of the user.
+ query: The query to match the memories on.
+
+ Returns:
+ A SearchMemoryResponse containing the matching memories.
+ """
+ if not self._tmpl_attrs.get("memory_service"):
+ self.set_up()
+ return await self._tmpl_attrs.get("memory_service").search_memory(
+ app_name=self._tmpl_attrs.get("app_name"),
+ user_id=user_id,
+ query=query,
+ )
+
  def register_operations(self) -> Dict[str, List[str]]:
  """Registers the operations of the ADK application."""
  return {
@@ -975,6 +1029,8 @@ def register_operations(self) -> Dict[str, List[str]]:
  "async_list_sessions",
  "async_create_session",
  "async_delete_session",
+ "async_add_session_to_memory",
+ "async_search_memory",
  ],
  "stream": ["stream_query", "streaming_agent_run_with_events"],
  "async_stream": ["async_stream_query"],