chore: Set agent_engine_id in the service constructor, also use the agent_engine_id field instead of overriding app_name in FastAPI endpoint

PiperOrigin-RevId: 770427903

Author: DeanChensj

src/google/adk/cli/fast_api.py

@@ -276,7 +276,6 @@ async def internal_lifespan(app: FastAPI):
  memory_service = InMemoryMemoryService()
 
  # Build the Session service
- agent_engine_id = ""
  if session_service_uri:
  if session_service_uri.startswith("agentengine://"):
  # Create vertex session service
@@ -285,8 +284,9 @@ async def internal_lifespan(app: FastAPI):
  raise click.ClickException("Agent engine id can not be empty.")
  envs.load_dotenv_for_agent("", agents_dir)
  session_service = VertexAiSessionService(
- os.environ["GOOGLE_CLOUD_PROJECT"],
- os.environ["GOOGLE_CLOUD_LOCATION"],
+ project=os.environ["GOOGLE_CLOUD_PROJECT"],
+ location=os.environ["GOOGLE_CLOUD_LOCATION"],
+ agent_engine_id=agent_engine_id,
  )
  else:
  session_service = DatabaseSessionService(db_url=session_service_uri)
@@ -357,8 +357,6 @@ def get_session_trace(session_id: str) -> Any:
  async def get_session(
  app_name: str, user_id: str, session_id: str
  ) -> Session:
- # Connect to managed session if agent_engine_id is set.
- app_name = agent_engine_id if agent_engine_id else app_name
  session = await session_service.get_session(
  app_name=app_name, user_id=user_id, session_id=session_id
  )
@@ -371,8 +369,6 @@ async def get_session(
  response_model_exclude_none=True,
  )
  async def list_sessions(app_name: str, user_id: str) -> list[Session]:
- # Connect to managed session if agent_engine_id is set.
- app_name = agent_engine_id if agent_engine_id else app_name
  list_sessions_response = await session_service.list_sessions(
  app_name=app_name, user_id=user_id
  )
@@ -393,8 +389,6 @@ async def create_session_with_id(
  session_id: str,
  state: Optional[dict[str, Any]] = None,
  ) -> Session:
- # Connect to managed session if agent_engine_id is set.
- app_name = agent_engine_id if agent_engine_id else app_name
  if (
  await session_service.get_session(
  app_name=app_name, user_id=user_id, session_id=session_id
@@ -419,8 +413,6 @@ async def create_session(
  user_id: str,
  state: Optional[dict[str, Any]] = None,
  ) -> Session:
- # Connect to managed session if agent_engine_id is set.
- app_name = agent_engine_id if agent_engine_id else app_name
  logger.info("New session created")
  return await session_service.create_session(
  app_name=app_name, user_id=user_id, state=state
@@ -660,8 +652,6 @@ def list_eval_results(app_name: str) -> list[str]:
 
  @app.delete("/apps/{app_name}/users/{user_id}/sessions/{session_id}")
  async def delete_session(app_name: str, user_id: str, session_id: str):
- # Connect to managed session if agent_engine_id is set.
- app_name = agent_engine_id if agent_engine_id else app_name
  await session_service.delete_session(
  app_name=app_name, user_id=user_id, session_id=session_id
  )
@@ -677,7 +667,6 @@ async def load_artifact(
  artifact_name: str,
  version: Optional[int] = Query(None),
  ) -> Optional[types.Part]:
- app_name = agent_engine_id if agent_engine_id else app_name
  artifact = await artifact_service.load_artifact(
  app_name=app_name,
  user_id=user_id,
@@ -700,7 +689,6 @@ async def load_artifact_version(
  artifact_name: str,
  version_id: int,
  ) -> Optional[types.Part]:
- app_name = agent_engine_id if agent_engine_id else app_name
  artifact = await artifact_service.load_artifact(
  app_name=app_name,
  user_id=user_id,
@@ -719,7 +707,6 @@ async def load_artifact_version(
  async def list_artifact_names(
  app_name: str, user_id: str, session_id: str
  ) -> list[str]:
- app_name = agent_engine_id if agent_engine_id else app_name
  return await artifact_service.list_artifact_keys(
  app_name=app_name, user_id=user_id, session_id=session_id
  )
@@ -731,7 +718,6 @@ async def list_artifact_names(
  async def list_artifact_versions(
  app_name: str, user_id: str, session_id: str, artifact_name: str
  ) -> list[int]:
- app_name = agent_engine_id if agent_engine_id else app_name
  return await artifact_service.list_versions(
  app_name=app_name,
  user_id=user_id,
@@ -745,7 +731,6 @@ async def list_artifact_versions(
  async def delete_artifact(
  app_name: str, user_id: str, session_id: str, artifact_name: str
  ):
- app_name = agent_engine_id if agent_engine_id else app_name
  await artifact_service.delete_artifact(
  app_name=app_name,
  user_id=user_id,
@@ -755,10 +740,8 @@ async def delete_artifact(
 
  @app.post("/run", response_model_exclude_none=True)
  async def agent_run(req: AgentRunRequest) -> list[Event]:
- # Connect to managed session if agent_engine_id is set.
- app_name = agent_engine_id if agent_engine_id else req.app_name
  session = await session_service.get_session(
- app_name=app_name, user_id=req.user_id, session_id=req.session_id
+ app_name=req.app_name, user_id=req.user_id, session_id=req.session_id
  )
  if not session:
  raise HTTPException(status_code=404, detail="Session not found")
@@ -776,11 +759,9 @@ async def agent_run(req: AgentRunRequest) -> list[Event]:
 
  @app.post("/run_sse")
  async def agent_run_sse(req: AgentRunRequest) -> StreamingResponse:
- # Connect to managed session if agent_engine_id is set.
- app_name = agent_engine_id if agent_engine_id else req.app_name
  # SSE endpoint
  session = await session_service.get_session(
- app_name=app_name, user_id=req.user_id, session_id=req.session_id
+ app_name=req.app_name, user_id=req.user_id, session_id=req.session_id
  )
  if not session:
  raise HTTPException(status_code=404, detail="Session not found")
@@ -818,8 +799,6 @@ async def event_generator():
  async def get_event_graph(
  app_name: str, user_id: str, session_id: str, event_id: str
  ):
- # Connect to managed session if agent_engine_id is set.
- app_name = agent_engine_id if agent_engine_id else app_name
  session = await session_service.get_session(
  app_name=app_name, user_id=user_id, session_id=session_id
  )
@@ -875,8 +854,6 @@ async def agent_live_run(
  ) -> None:
  await websocket.accept()
 
- # Connect to managed session if agent_engine_id is set.
- app_name = agent_engine_id if agent_engine_id else app_name
  session = await session_service.get_session(
  app_name=app_name, user_id=user_id, session_id=session_id
  )
@@ -940,7 +917,7 @@ async def _get_runner_async(app_name: str) -> Runner:
  return runner_dict[app_name]
  root_agent = agent_loader.load_agent(app_name)
  runner = Runner(
- app_name=agent_engine_id if agent_engine_id else app_name,
+ app_name=app_name,
  agent=root_agent,
  artifact_service=artifact_service,
  session_service=session_service,

src/google/adk/sessions/vertex_ai_session_service.py

@@ -22,7 +22,6 @@
 import urllib.parse
 
 from dateutil import parser
-from google.genai import types
 from typing_extensions import override
 
 from google import genai
@@ -40,15 +39,27 @@
 
 
 class VertexAiSessionService(BaseSessionService):
- """Connects to the managed Vertex AI Session Service."""
+ """Connects to the Vertex AI Agent Engine Session Service using GenAI API client.
+
+ https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/sessions/overview
+ """
 
  def __init__(
  self,
- project: str = None,
- location: str = None,
+ project: Optional[str] = None,
+ location: Optional[str] = None,
+ agent_engine_id: Optional[str] = None,
  ):
- self.project = project
- self.location = location
+ """Initializes the VertexAiSessionService.
+
+ Args:
+ project: The project id of the project to use.
+ location: The location of the project to use.
+ agent_engine_id: The resource ID of the agent engine to use.
+ """
+ self._project = project
+ self._location = location
+ self._agent_engine_id = agent_engine_id
 
  @override
  async def create_session(
@@ -64,14 +75,13 @@ async def create_session(
  'User-provided Session id is not supported for'
  ' VertexAISessionService.'
  )
-
- reasoning_engine_id = _parse_reasoning_engine_id(app_name)
+ reasoning_engine_id = self._get_reasoning_engine_id(app_name)
+ api_client = self._get_api_client()
 
  session_json_dict = {'user_id': user_id}
  if state:
  session_json_dict['session_state'] = state
 
- api_client = _get_api_client(self.project, self.location)
  api_response = await api_client.async_request(
  http_method='POST',
  path=f'reasoningEngines/{reasoning_engine_id}/sessions',
@@ -130,10 +140,10 @@ async def get_session(
  session_id: str,
  config: Optional[GetSessionConfig] = None,
  ) -> Optional[Session]:
- reasoning_engine_id = _parse_reasoning_engine_id(app_name)
+ reasoning_engine_id = self._get_reasoning_engine_id(app_name)
+ api_client = self._get_api_client()
 
  # Get session resource
- api_client = _get_api_client(self.project, self.location)
  get_session_api_response = await api_client.async_request(
  http_method='GET',
  path=f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}',
@@ -203,14 +213,14 @@ async def get_session(
  async def list_sessions(
  self, *, app_name: str, user_id: str
  ) -> ListSessionsResponse:
- reasoning_engine_id = _parse_reasoning_engine_id(app_name)
+ reasoning_engine_id = self._get_reasoning_engine_id(app_name)
+ api_client = self._get_api_client()
 
  path = f'reasoningEngines/{reasoning_engine_id}/sessions'
  if user_id:
  parsed_user_id = urllib.parse.quote(f'''"{user_id}"''', safe='')
  path = path + f'?filter=user_id={parsed_user_id}'
 
- api_client = _get_api_client(self.project, self.location)
  api_response = await api_client.async_request(
  http_method='GET',
  path=path,
@@ -236,8 +246,9 @@ async def list_sessions(
  async def delete_session(
  self, *, app_name: str, user_id: str, session_id: str
  ) -> None:
- reasoning_engine_id = _parse_reasoning_engine_id(app_name)
- api_client = _get_api_client(self.project, self.location)
+ reasoning_engine_id = self._get_reasoning_engine_id(app_name)
+ api_client = self._get_api_client()
+
  try:
  await api_client.async_request(
  http_method='DELETE',
@@ -253,24 +264,43 @@ async def append_event(self, session: Session, event: Event) -> Event:
  # Update the in-memory session.
  await super().append_event(session=session, event=event)
 
- reasoning_engine_id = _parse_reasoning_engine_id(session.app_name)
- api_client = _get_api_client(self.project, self.location)
+ reasoning_engine_id = self._get_reasoning_engine_id(session.app_name)
+ api_client = self._get_api_client()
  await api_client.async_request(
  http_method='POST',
  path=f'reasoningEngines/{reasoning_engine_id}/sessions/{session.id}:appendEvent',
  request_dict=_convert_event_to_json(event),
  )
  return event
 
+ def _get_reasoning_engine_id(self, app_name: str):
+ if self._agent_engine_id:
+ return self._agent_engine_id
 
-def _get_api_client(project: str, location: str):
- """Instantiates an API client for the given project and location.
+ if app_name.isdigit():
+  return app_name
 
- It needs to be instantiated inside each request so that the event loop
- management.
- """
- client = genai.Client(vertexai=True, project=project, location=location)
- return client._api_client
+ pattern = r'^projects/([a-zA-Z0-9-_]+)/locations/([a-zA-Z0-9-_]+)/reasoningEngines/(\d+)$'
+ match = re.fullmatch(pattern, app_name)
+
+ if not bool(match):
+ raise ValueError(
+ f'App name {app_name} is not valid. It should either be the full'
+ ' ReasoningEngine resource name, or the reasoning engine id.'
+ )
+
+ return match.groups()[-1]
+
+ def _get_api_client(self):
+ """Instantiates an API client for the given project and location.
+
+ It needs to be instantiated inside each request so that the event loop
+ management can be properly propagated.
+ """
+ client = genai.Client(
+ vertexai=True, project=self._project, location=self._location
+ )
+ return client._api_client
 
 
 def _convert_event_to_json(event: Event) -> Dict[str, Any]:
@@ -366,19 +396,3 @@ def _from_api_event(api_event: Dict[str, Any]) -> Event:
  )
 
  return event
-
-
-def _parse_reasoning_engine_id(app_name: str):
- if app_name.isdigit():
- return app_name
-
- pattern = r'^projects/([a-zA-Z0-9-_]+)/locations/([a-zA-Z0-9-_]+)/reasoningEngines/(\d+)$'
- match = re.fullmatch(pattern, app_name)
-
- if not bool(match):
- raise ValueError(
- f'App name {app_name} is not valid. It should either be the full'
- ' ReasoningEngine resource name, or the reasoning engine id.'
- )
-
- return match.groups()[-1]

tests/unittests/sessions/test_vertex_ai_session_service.py

@@ -245,8 +245,14 @@ async def async_request(
  raise ValueError(f'Unsupported http method: {http_method}')
 
 
-def mock_vertex_ai_session_service():
+def mock_vertex_ai_session_service(agent_engine_id: Optional[str] = None):
  """Creates a mock Vertex AI Session service for testing."""
+ if agent_engine_id:
+ return VertexAiSessionService(
+ project='test-project',
+ location='test-location',
+ agent_engine_id=agent_engine_id,
+ )
  return VertexAiSessionService(
  project='test-project', location='test-location'
  )
@@ -265,16 +271,20 @@ def mock_get_api_client():
  '2': (MOCK_EVENT_JSON_2, 'my_token'),
  }
  with mock.patch(
- 'google.adk.sessions.vertex_ai_session_service._get_api_client',
+ 'google.adk.sessions.vertex_ai_session_service.VertexAiSessionService._get_api_client',
  return_value=api_client,
  ):
  yield
 
 
 @pytest.mark.asyncio
 @pytest.mark.usefixtures('mock_get_api_client')
-async def test_get_empty_session():
- session_service = mock_vertex_ai_session_service()
+@pytest.mark.parametrize('agent_engine_id', [None, '123'])
+async def test_get_empty_session(agent_engine_id):
+ if agent_engine_id:
+ session_service = mock_vertex_ai_session_service(agent_engine_id)
+ else:
+ session_service = mock_vertex_ai_session_service()
  with pytest.raises(ValueError) as excinfo:
  await session_service.get_session(
  app_name='123', user_id='user', session_id='0'