feat: Add Vertex Express mode compatibility for VertexAiSessionService

PiperOrigin-RevId: 775317848

src/google/adk/sessions/vertex_ai_session_service.py

@@ -16,6 +16,7 @@ from __future__ import annotations
 import asyncio
 import json
 import logging
+import os
 import re
 from typing import Any
 from typing import Dict
@@ -23,6 +24,7 @@ from typing import Optional
 import urllib.parse
 
 from dateutil import parser
+from google.genai.errors import ClientError
 from typing_extensions import override
 
 from google import genai
@@ -95,8 +97,30 @@ class VertexAiSessionService(BaseSessionService):
     operation_id = api_response['name'].split('/')[-1]
 
     max_retry_attempt = 5
+
+    if _is_vertex_express_mode(self._project, self._location):
+      # Express mode doesn't support LRO, so we need to poll
+      # the session resource.
+      # TODO: remove this once LRO polling is supported in Express mode.
+      for i in range(max_retry_attempt):
+        try:
+          await api_client.async_request(
+              http_method='GET',
+              path=(
+                  f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}'
+              ),
+              request_dict={},
+          )
+          break
+        except ClientError as e:
+          logger.info('Polling for session %s: %s', session_id, e)
+          # Add slight exponential backoff to avoid excessive polling.
+          await asyncio.sleep(1 + 0.5 * i)
+      else:
+        raise TimeoutError('Session creation failed.')
+    else:
       lro_response = None
-    while max_retry_attempt >= 0:
+      for _ in range(max_retry_attempt):
         lro_response = await api_client.async_request(
             http_method='GET',
             path=f'operations/{operation_id}',
@@ -108,7 +132,6 @@ class VertexAiSessionService(BaseSessionService):
           break
 
         await asyncio.sleep(1)
-      max_retry_attempt -= 1
 
       if lro_response is None or not lro_response.get('done', None):
         raise TimeoutError(
@@ -312,6 +335,18 @@ class VertexAiSessionService(BaseSessionService):
     return client._api_client
 
 
+def _is_vertex_express_mode(
+    project: Optional[str], location: Optional[str]
+) -> bool:
+  """Check if Vertex AI and API key are both enabled replacing project and location, meaning the user is using the Vertex Express Mode."""
+  return (
+      os.environ.get('GOOGLE_GENAI_USE_VERTEXAI', '0').lower() in ['true', '1']
+      and os.environ.get('GOOGLE_API_KEY', None) is not None
+      and project is None
+      and location is None
+  )
+
+
 def _convert_api_response(api_response):
   """Converts the API response to a JSON object based on the type."""
   if hasattr(api_response, 'body'):