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

PiperOrigin-RevId: 803656127

Author: vertex-sdk-bot

tests/unit/vertexai/genai/test_agent_engines.py

@@ -149,6 +149,20 @@ async def custom_async_stream_method(
         for chunk in _TEST_AGENT_ENGINE_STREAM_QUERY_RESPONSE:
             yield chunk
 
+    async def bidi_stream_query(self, input_queue: asyncio.Queue) -> AsyncIterable[Any]:
+        """Runs the bidi stream engine."""
+        while True:
+            chunk = await input_queue.get()
+            yield chunk
+
+    async def custom_bidi_stream_method(
+        self, input_queue: asyncio.Queue
+    ) -> AsyncIterable[Any]:
+        """Runs the async bidi stream engine."""
+        while True:
+            chunk = await input_queue.get()
+            yield chunk
+
     def clone(self):
         return self
 
@@ -170,6 +184,10 @@ def register_operations(self) -> Dict[str, List[str]]:
                 _TEST_DEFAULT_ASYNC_STREAM_METHOD_NAME,
                 _TEST_CUSTOM_ASYNC_STREAM_METHOD_NAME,
             ],
+            _TEST_BIDI_STREAM_API_MODE: [
+                _TEST_DEFAULT_BIDI_STREAM_METHOD_NAME,
+                _TEST_CUSTOM_BIDI_STREAM_METHOD_NAME,
+            ],
         }
 
 
@@ -323,21 +341,27 @@ def register_operations(self) -> Dict[str, List[str]]:
 _TEST_ASYNC_API_MODE = _agent_engines_utils._ASYNC_API_MODE
 _TEST_STREAM_API_MODE = _agent_engines_utils._STREAM_API_MODE
 _TEST_ASYNC_STREAM_API_MODE = _agent_engines_utils._ASYNC_STREAM_API_MODE
+_TEST_BIDI_STREAM_API_MODE = _agent_engines_utils._BIDI_STREAM_API_MODE
 _TEST_DEFAULT_METHOD_NAME = _agent_engines_utils._DEFAULT_METHOD_NAME
 _TEST_DEFAULT_ASYNC_METHOD_NAME = _agent_engines_utils._DEFAULT_ASYNC_METHOD_NAME
 _TEST_DEFAULT_STREAM_METHOD_NAME = _agent_engines_utils._DEFAULT_STREAM_METHOD_NAME
 _TEST_DEFAULT_ASYNC_STREAM_METHOD_NAME = (
     _agent_engines_utils._DEFAULT_ASYNC_STREAM_METHOD_NAME
 )
+_TEST_DEFAULT_BIDI_STREAM_METHOD_NAME = (
+    _agent_engines_utils._DEFAULT_BIDI_STREAM_METHOD_NAME
+)
 _TEST_CAPITALIZE_ENGINE_METHOD_DOCSTRING = "Runs the engine."
 _TEST_STREAM_METHOD_DOCSTRING = "Runs the stream engine."
 _TEST_ASYNC_STREAM_METHOD_DOCSTRING = "Runs the async stream engine."
+_TEST_BIDI_STREAM_METHOD_DOCSTRING = "Runs the bidi stream engine."
 _TEST_MODE_KEY_IN_SCHEMA = _agent_engines_utils._MODE_KEY_IN_SCHEMA
 _TEST_METHOD_NAME_KEY_IN_SCHEMA = _agent_engines_utils._METHOD_NAME_KEY_IN_SCHEMA
 _TEST_CUSTOM_METHOD_NAME = "custom_method"
 _TEST_CUSTOM_ASYNC_METHOD_NAME = "custom_async_method"
 _TEST_CUSTOM_STREAM_METHOD_NAME = "custom_stream_method"
 _TEST_CUSTOM_ASYNC_STREAM_METHOD_NAME = "custom_async_stream_method"
+_TEST_CUSTOM_BIDI_STREAM_METHOD_NAME = "custom_bidi_stream_method"
 _TEST_CUSTOM_METHOD_DEFAULT_DOCSTRING = """
     Runs the Agent Engine to serve the user request.
 
@@ -579,6 +603,22 @@ def register_operations(self) -> Dict[str, List[str]]:
 _TEST_AGENT_ENGINE_CUSTOM_ASYNC_STREAM_QUERY_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = (
     _TEST_ASYNC_STREAM_API_MODE
 )
+_TEST_AGENT_ENGINE_BIDI_STREAM_QUERY_SCHEMA = _agent_engines_utils._generate_schema(
+    OperationRegistrableEngine().bidi_stream_query,
+    schema_name=_TEST_DEFAULT_BIDI_STREAM_METHOD_NAME,
+)
+_TEST_AGENT_ENGINE_BIDI_STREAM_QUERY_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = (
+    _TEST_BIDI_STREAM_API_MODE
+)
+_TEST_AGENT_ENGINE_CUSTOM_BIDI_STREAM_QUERY_SCHEMA = (
+    _agent_engines_utils._generate_schema(
+        OperationRegistrableEngine().custom_bidi_stream_method,
+        schema_name=_TEST_CUSTOM_BIDI_STREAM_METHOD_NAME,
+    )
+)
+_TEST_AGENT_ENGINE_CUSTOM_BIDI_STREAM_QUERY_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = (
+    _TEST_BIDI_STREAM_API_MODE
+)
 _TEST_OPERATION_REGISTRABLE_SCHEMAS = [
     _TEST_AGENT_ENGINE_QUERY_SCHEMA,
     _TEST_AGENT_ENGINE_CUSTOM_METHOD_SCHEMA,
@@ -588,6 +628,8 @@ def register_operations(self) -> Dict[str, List[str]]:
     _TEST_AGENT_ENGINE_CUSTOM_STREAM_QUERY_SCHEMA,
     _TEST_AGENT_ENGINE_ASYNC_STREAM_QUERY_SCHEMA,
     _TEST_AGENT_ENGINE_CUSTOM_ASYNC_STREAM_QUERY_SCHEMA,
+    _TEST_AGENT_ENGINE_BIDI_STREAM_QUERY_SCHEMA,
+    _TEST_AGENT_ENGINE_CUSTOM_BIDI_STREAM_QUERY_SCHEMA,
 ]
 _TEST_OPERATION_NOT_REGISTERED_SCHEMAS = [
     _TEST_AGENT_ENGINE_CUSTOM_METHOD_SCHEMA,
@@ -1873,6 +1915,20 @@ async def consume():
                         ),
                         _TEST_ASYNC_STREAM_API_MODE,
                     ),
+                    (
+                        _agent_engines_utils._generate_schema(
+                            OperationRegistrableEngine().bidi_stream_query,
+                            schema_name=_TEST_DEFAULT_BIDI_STREAM_METHOD_NAME,
+                        ),
+                        _TEST_BIDI_STREAM_API_MODE,
+                    ),
+                    (
+                        _agent_engines_utils._generate_schema(
+                            OperationRegistrableEngine().custom_bidi_stream_method,
+                            schema_name=_TEST_CUSTOM_BIDI_STREAM_METHOD_NAME,
+                        ),
+                        _TEST_BIDI_STREAM_API_MODE,
+                    ),
                 ],
             ),
             (

vertexai/_genai/_agent_engines_utils.py

@@ -15,6 +15,7 @@
 """Utility functions for agent engines."""
 
 import abc
+import asyncio
 from importlib import metadata as importlib_metadata
 import inspect
 import io
@@ -108,6 +109,7 @@
 _AGENT_FRAMEWORK_ATTR = "agent_framework"
 _ASYNC_API_MODE = "async"
 _ASYNC_STREAM_API_MODE = "async_stream"
+_BIDI_STREAM_API_MODE = "bidi_stream"
 _BASE_MODULES = set(_BUILTIN_MODULE_NAMES + tuple(_STDLIB_MODULE_NAMES))
 _BLOB_FILENAME = "agent_engine.pkl"
 _DEFAULT_AGENT_FRAMEWORK = "custom"
@@ -132,6 +134,7 @@
 _DEFAULT_STREAM_METHOD_RETURN_TYPE = "Iterable[Any]"
 _DEFAULT_REQUIRED_PACKAGES = frozenset(["cloudpickle", "pydantic"])
 _DEFAULT_STREAM_METHOD_NAME = "stream_query"
+_DEFAULT_BIDI_STREAM_METHOD_NAME = "bidi_stream_query"
 _EXTRA_PACKAGES_FILE = "dependencies.tar.gz"
 _FAILED_TO_REGISTER_API_METHODS_WARNING_TEMPLATE = (
     "Failed to register API methods. Please follow the guide to "
@@ -202,6 +205,15 @@ def stream_query(self, **kwargs) -> Iterator[Any]:  # type: ignore[no-untyped-de
         """Stream responses to serve the user query."""
 
 
+@typing.runtime_checkable
+class BidiStreamQueryable(Protocol):
+    """Protocol for Agent Engines that can stream requests and responses."""
+
+    @abc.abstractmethod
+    async def bidi_stream_query(self, input_queue: asyncio.Queue) -> AsyncIterator[Any]:
+        """Stream requests and responses to serve the user queries."""
+
+
 @typing.runtime_checkable
 class Cloneable(Protocol):
     """Protocol for Agent Engines that can be cloned."""
@@ -234,6 +246,7 @@ def register_operations(self, **kwargs) -> Dict[str, Sequence[str]]:
     OperationRegistrable,
     Queryable,
     StreamQueryable,
+    BidiStreamQueryable,
 ]
 
 
@@ -557,6 +570,9 @@ def _generate_schema(
             inspect.Parameter.KEYWORD_ONLY,
             inspect.Parameter.POSITIONAL_ONLY,
         )
+        # For a bidi endpoint, it requires an asyncio.Queue as the input, but
+        # it is not JSON serializable. We hence exclude it from the schema.
+        and param.annotation != asyncio.Queue
     }
     parameters = pydantic.create_model(f.__name__, **fields_dict).schema()
     # Postprocessing
@@ -656,6 +672,8 @@ def _get_registered_operations(
         operations[_STREAM_API_MODE] = [_DEFAULT_STREAM_METHOD_NAME]
     if isinstance(agent, AsyncStreamQueryable):
         operations[_ASYNC_STREAM_API_MODE] = [_DEFAULT_ASYNC_STREAM_METHOD_NAME]
+    if isinstance(agent, BidiStreamQueryable):
+        operations[_BIDI_STREAM_API_MODE] = [_DEFAULT_BIDI_STREAM_METHOD_NAME]
     return operations
 
 
@@ -839,6 +857,10 @@ def _register_api_methods_or_raise(
                 f" contain an `{_MODE_KEY_IN_SCHEMA}` field."
             )
         api_mode = operation_schema.get(_MODE_KEY_IN_SCHEMA)
+        # For bidi stream api mode, we don't need to wrap the operation.
+        if api_mode == _BIDI_STREAM_API_MODE:
+            continue
+
         if _METHOD_NAME_KEY_IN_SCHEMA not in operation_schema:
             raise ValueError(
                 f"Operation schema {operation_schema} does not"
@@ -1212,6 +1234,7 @@ def _validate_agent_or_raise(
     * a callable method named `query`
     * a callable method named `stream_query`
     * a callable method named `async_stream_query`
+    * a callable method named `bidi_stream_query`
     * a callable method named `register_operations`
 
     Args:
@@ -1246,6 +1269,9 @@ def _validate_agent_or_raise(
     is_async_stream_queryable = isinstance(agent, AsyncStreamQueryable) and callable(
         agent.async_stream_query
     )
+    is_bidi_stream_queryable = isinstance(agent, BidiStreamQueryable) and callable(
+        agent.bidi_stream_query
+    )
     is_operation_registrable = isinstance(agent, OperationRegistrable) and callable(
         agent.register_operations
     )
@@ -1255,12 +1281,13 @@ def _validate_agent_or_raise(
         or is_async_queryable
         or is_stream_queryable
         or is_operation_registrable
+        or is_bidi_stream_queryable
         or is_async_stream_queryable
     ):
         raise TypeError(
             "agent_engine has none of the following callable methods: "
-            "`query`, `async_query`, `stream_query`, `async_stream_query` or "
-            "`register_operations`."
+            "`query`, `async_query`, `stream_query`, `async_stream_query`, "
+            "`bidi_stream_query`, or `register_operations`."
         )
 
     if is_queryable:
@@ -1299,6 +1326,15 @@ def _validate_agent_or_raise(
                 " missing `self` argument in the agent.async_stream_query method."
             ) from err
 
+    if is_bidi_stream_queryable:
+        try:
+            inspect.signature(getattr(agent, "bidi_stream_query"))
+        except ValueError as err:
+            raise ValueError(
+                "Invalid bidi_stream_query signature. This might be due to a "
+                " missing `self` argument in the agent.bidi_stream_query method."
+            ) from err
+
     if is_operation_registrable:
         try:
             inspect.signature(getattr(agent, "register_operations"))

vertexai/_genai/agent_engines.py

@@ -1220,6 +1220,18 @@ def _create_config(
                 _agent_engines_utils._to_dict(class_method)
                 for class_method in class_methods
             ]
+            # Set the agent_server_mode to EXPERIMENTAL if the agent has a
+            # bidi_stream method.
+            for class_method in class_methods:
+                if class_method["api_mode"] == "bidi_stream":
+                    if not agent_engine_spec.get("deployment_spec"):
+                        agent_engine_spec["deployment_spec"] = (
+                            types.ReasoningEngineSpecDeploymentSpecDict()
+                        )
+                    agent_engine_spec["deployment_spec"][
+                        "agent_server_mode"
+                    ] = types.AgentServerMode.EXPERIMENTAL
+                    break
             update_masks.append("spec.class_methods")
             agent_engine_spec["agent_framework"] = (
                 _agent_engines_utils._get_agent_framework(agent=agent)