feat: GenAI SDK client - Add live/bidi streaming support for Agent Engine

PiperOrigin-RevId: 799712899

Author: vertex-sdk-bot

tests/unit/vertexai/genai/test_genai_client.py

@@ -56,6 +56,12 @@ async def test_async_client(self):
  test_client = vertexai.Client(project=_TEST_PROJECT, location=_TEST_LOCATION)
  assert isinstance(test_client.aio, vertexai._genai.client.AsyncClient)
 
+ @pytest.mark.usefixtures("google_auth_mock")
+ def test_live_client(self):
+ test_client = vertexai.Client(project=_TEST_PROJECT, location=_TEST_LOCATION)
+ test_async_client = test_client.aio
+ assert isinstance(test_async_client.live, vertexai._genai.live.AsyncLive)
+
  @pytest.mark.usefixtures("google_auth_mock")
  def test_types(self):
  assert vertexai.types is not None

tests/unit/vertexai/genai/test_live_agent_engines.py

@@ -0,0 +1,112 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import importlib
+import json
+from unittest import mock
+
+import google.auth
+import google.auth.credentials
+from google.cloud import aiplatform
+import vertexai
+from google.cloud.aiplatform import initializer as aiplatform_initializer
+from vertexai._genai import live_agent_engines
+import pytest
+
+
+_TEST_PROJECT = "test-project"
+_TEST_LOCATION = "us-central1"
+pytestmark = pytest.mark.usefixtures("google_auth_mock")
+
+
+class TestLiveAgentEngines:
+ """Unit tests for the GenAI client."""
+
+ def setup_method(self):
+ importlib.reload(aiplatform_initializer)
+ importlib.reload(aiplatform)
+ importlib.reload(vertexai)
+ vertexai.init(
+ project=_TEST_PROJECT,
+ location=_TEST_LOCATION,
+ )
+
+ @pytest.mark.asyncio
+ @pytest.mark.usefixtures("google_auth_mock")
+ @mock.patch.object(live_agent_engines, "ws_connect")
+ @mock.patch.object(google.auth, "default")
+ async def test_async_live_agent_engines_connect(
+ self, mock_auth_default, mock_ws_connect
+ ):
+ """Tests the AsyncLiveAgentEngines.connect method, as well as the AsyncLiveAgentEngineSession methods."""
+ # Mock credentials to avoid refresh issues
+ mock_creds = mock.Mock(spec=google.auth.credentials.Credentials)
+ mock_creds.token = "test-token"
+ mock_creds.valid = True
+ mock_auth_default.return_value = (mock_creds, None)
+
+ test_client = vertexai.Client(project=_TEST_PROJECT, location=_TEST_LOCATION)
+ mock_ws = mock.AsyncMock()
+ mock_ws_connect.return_value.__aenter__.return_value = mock_ws
+
+ mock_ws.recv.side_effect = [
+ json.dumps({"output": "HELLO"}).encode("utf-8"),
+ json.dumps({"output": "WORLD"}).encode("utf-8"),
+ ]
+
+ async with test_client.aio.live.agent_engines.connect(
+ agent_engine="test-agent-engine",
+ config={"class_method": "bidi_stream_query", "input": {"query": "hello"}},
+ ) as session:
+ assert session is not None
+
+ # Send additional messages
+ await session.send({"query": "world"})
+
+ # Receive responses
+ responses = []
+ response = await session.receive()
+ responses.append(response)
+ response = await session.receive()
+ responses.append(response)
+
+ await session.close()
+
+ assert responses == [{"output": "HELLO"}, {"output": "WORLD"}]
+
+ mock_ws.send.assert_has_calls(
+ [
+ mock.call(
+ json.dumps(
+ {
+ "setup": {
+ "name": (
+ f"projects/{_TEST_PROJECT}/locations/"
+ f"{_TEST_LOCATION}/reasoningEngines/"
+ "test-agent-engine"
+ ),
+ "class_method": "bidi_stream_query",
+ "input": {"query": "hello"},
+ }
+ }
+ )
+ ),
+ mock.call(json.dumps({"bidi_stream_input": {"query": "world"}})),
+ ]
+ )
+ mock_ws.close.assert_called_once()
+ mock_auth_default.assert_called_once_with(
+ scopes=["https://www.googleapis.com/auth/cloud-platform"]
+ )
+ mock_creds.refresh.assert_not_called()

vertexai/_genai/client.py

@@ -22,6 +22,7 @@
 from google.genai import _common
 from google.genai import client as genai_client
 from google.genai import types
+from . import live
 
 
 _GENAI_MODULES_TELEMETRY_HEADER = "vertex-genai-modules"
@@ -45,15 +46,23 @@ def _add_tracking_headers(headers: dict[str, str]) -> None:
 
 
 class AsyncClient:
-
  """Async Client for the GenAI SDK."""
 
  def __init__(self, api_client: genai_client.Client):
  self._api_client = api_client
+ self._live = live.AsyncLive(self._api_client)
  self._evals = None
  self._agent_engines = None
  self._prompt_optimizer = None
 
+ @property
+ @_common.experimental_warning(
+ "The Vertex SDK GenAI live module is experimental, and may change in future "
+ "versions."
+ )
+ def live(self) -> live.AsyncLive:
+ return self._live
+
  @property
  @_common.experimental_warning(
  "The Vertex SDK GenAI evals module is experimental, and may change in future "

vertexai/_genai/live.py

@@ -0,0 +1,57 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""[Preview] Live API client."""
+
+import importlib
+import logging
+
+
+from google.genai import _api_module
+from google.genai import _common
+from google.genai._api_client import BaseApiClient
+
+logger = logging.getLogger("google_genai.live")
+
+
+class AsyncLive(_api_module.BaseModule):
+ """[Preview] AsyncLive."""
+
+ def __init__(self, api_client: BaseApiClient):
+ super().__init__(api_client)
+ self._agent_engines = None
+
+ @property
+ @_common.experimental_warning(
+ "The Vertex SDK GenAI agent engines module is experimental, "
+ "and may change in future versions."
+ )
+ def agent_engines(self):
+ if self._agent_engines is None:
+ try:
+ # We need to lazy load the live_agent_engines module to handle
+ # the possibility of ImportError when dependencies are not
+ # installed.
+ self._agent_engines = importlib.import_module(
+ ".live_agent_engines",
+ __package__,
+ )
+ except ImportError as e:
+ raise ImportError(
+ "The 'agent_engines' module requires 'additional packages'. "
+ "Please install them using pip install "
+ "google-cloud-aiplatform[agent_engines]"
+ ) from e
+ return self._agent_engines.AsyncLiveAgentEngines(self._api_client)

vertexai/_genai/live_agent_engines.py

@@ -0,0 +1,179 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Live AgentEngine API client."""
+
+import contextlib
+import json
+from typing import Any, AsyncIterator, Dict, Optional
+import google.auth
+
+from google.genai import _api_module
+from .types import QueryAgentEngineConfig, QueryAgentEngineConfigOrDict
+
+
+try:
+ from websockets.asyncio.client import ClientConnection
+ from websockets.asyncio.client import connect as ws_connect
+except ModuleNotFoundError:
+ # This try/except is for TAP, mypy complains about it which is why we have the type: ignore
+ from websockets.client import ClientConnection # type: ignore
+ from websockets.client import connect as ws_connect # type: ignore
+
+
+class AsyncLiveAgentEngineSession:
+ """AsyncLiveAgentEngineSession."""
+
+ def __init__(self, websocket: ClientConnection):
+ self._ws = websocket
+
+ async def send(self, query_input: Dict[str, Any]) -> None:
+ """Send a query input to the Agent.
+
+ Args:
+ query_input: A JSON serializable Python Dict to be send to the Agent.
+ """
+
+ try:
+ json_request = json.dumps({"bidi_stream_input": query_input})
+ except json.JSONEncoderError as exc:
+ raise ValueError(
+ "Failed to encode query input to JSON in live_agent_engines: "
+ f"{str(query_input)}"
+ ) from exc
+ await self._ws.send(json_request)
+
+ async def receive(self) -> Dict[str, Any]:
+ """Receive one response from the Agent.
+
+ Returns:
+ A response from the Agent.
+
+ Raises:
+ websockets.exceptions.ConnectionClosed: If the connection is closed.
+ """
+
+ response = await self._ws.recv()
+ try:
+ return json.loads(response)
+ except json.decoder.JSONDecodeError as exc:
+ raise ValueError(
+ "Failed to parse response to JSON in live_agent_engines: "
+ f"{str(response)}"
+ ) from exc
+
+ async def close(self) -> None:
+ """Close the connection."""
+ await self._ws.close()
+
+
+class AsyncLiveAgentEngines(_api_module.BaseModule):
+ """AsyncLiveAgentEngines.
+
+ Example usage:
+
+ .. code-block:: python
+
+ from pathlib import Path
+
+ from google import genai
+ from google.genai import types
+
+ class MyAgentEngine(client):
+ def bidi_stream_query(self, input_queue: asyncio.Queue):
+ while True:
+ input = await input_queue.get()
+ yield {"output": f"Agent received {input}!"}
+
+ client = vertexai.Client(project="my-project", location="us-central1")
+ agent_engine = client.agent_engines.create(agent)
+
+ async with client.aio.live.agent_engines.connect(
+ agent_engine=agent_engine.api_resource.name,
+ setup={"class_method": "bidi_stream_query"},
+ ) as session:
+ await session.send(input={"input": "Hello world"})
+
+ response = await session.receive()
+ # {"output": "Agent received Hello world!"}
+ ...
+ """
+
+ @contextlib.asynccontextmanager
+ async def connect(
+ self,
+ *,
+ agent_engine: str,
+ config: Optional[QueryAgentEngineConfigOrDict] = None,
+ ) -> AsyncIterator[AsyncLiveAgentEngineSession]:
+ """Connect to the agent deployed to Agent Engine in a live (bidirectional streaming) session.
+
+ Args:
+ agent_engine: The resource name of the Agent Engine to use for the
+ live session.
+ config: The optional configuration for starting the live Agent Engine
+ session. Custom class_method and an optional initial input could be
+ provided. If no class_method is provided, the default class_method
+ "bidi_stream_query" will be used by the Agent Engine.
+
+ Yields:
+ An AsyncLiveAgentEngineSession object.
+ """
+ if isinstance(config, dict):
+ config = QueryAgentEngineConfig(**config)
+
+ agent_engine_resource_name = agent_engine
+ if not agent_engine_resource_name.startswith("projects/"):
+ agent_engine_resource_name = f"projects/{self._api_client.project}/locations/{self._api_client.location}/reasoningEngines/{agent_engine}"
+ request_dict = {"setup": {"name": agent_engine_resource_name}}
+ if config.class_method:
+ request_dict["setup"]["class_method"] = config.class_method
+ if config.input:
+ request_dict["setup"]["input"] = config.input
+
+ request = json.dumps(request_dict)
+
+ if not self._api_client._credentials:
+ # Get bearer token through Application Default Credentials.
+ creds, _ = google.auth.default( # type: ignore
+ scopes=["https://www.googleapis.com/auth/cloud-platform"]
+ )
+ else:
+ creds = self._api_client._credentials
+ # creds.valid is False, and creds.token is None
+ # Need to refresh credentials to populate those
+ if not (creds.token and creds.valid):
+ auth_req = google.auth.transport.requests.Request() # type: ignore
+ creds.refresh(auth_req)
+ bearer_token = creds.token
+
+ original_headers = self._api_client._http_options.headers
+ headers = original_headers.copy() if original_headers is not None else {}
+ headers["Authorization"] = f"Bearer {bearer_token}"
+
+ base_url = self._api_client._websocket_base_url()
+ if isinstance(base_url, bytes):
+ base_url = base_url.decode("utf-8")
+ uri = (
+ f"{base_url}/ws/google.cloud.aiplatform."
+ f"{self._api_client._http_options.api_version}"
+ ".ReasoningEngineExecutionService/BidiQueryReasoningEngine"
+ )
+
+ async with ws_connect(
+ uri, additional_headers=headers, **self._api_client._websocket_ssl_ctx
+ ) as ws:
+ await ws.send(request)
+ yield AsyncLiveAgentEngineSession(websocket=ws)