feat: Support bidi stream query in agent engines and ADK template.

PiperOrigin-RevId: 800547989

Author: vertex-sdk-bot

tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+import asyncio
 import base64
 import importlib
 import json
@@ -180,6 +181,34 @@ async def run_async(self, *args, **kwargs):
  }
  )
 
+ async def run_live(self, *args, **kwargs):
+ from google.adk.events import event
+
+ yield event.Event(
+ **{
+ "author": "currency_exchange_agent",
+ "content": {
+ "parts": [
+ {
+ "thought_signature": b"test_signature",
+ "function_call": {
+ "args": {
+ "currency_date": "2025-04-03",
+ "currency_from": "USD",
+ "currency_to": "SEK",
+ },
+ "id": "af-c5a57692-9177-4091-a3df-098f834ee849",
+ "name": "get_exchange_rate",
+ },
+ }
+ ],
+ "role": "model",
+ },
+ "id": "9aaItGK9",
+ "invocation_id": "e-6543c213-6417-484b-9551-b67915d1d5f7",
+ }
+ )
+
 
 @pytest.mark.usefixtures("google_auth_mock")
 class TestAdkApp:
@@ -355,6 +384,29 @@ def test_streaming_agent_run_with_events(self):
  events = list(app.streaming_agent_run_with_events(request_json=request_json))
  assert len(events) == 1
 
+ @pytest.mark.asyncio
+ async def test_async_bidi_stream_query(self):
+ app = reasoning_engines.AdkApp(
+ agent=Agent(name=_TEST_AGENT_NAME, model=_TEST_MODEL)
+ )
+ assert app._tmpl_attrs.get("runner") is None
+ app.set_up()
+ app._tmpl_attrs["runner"] = _MockRunner()
+ request_queue = asyncio.Queue()
+ request_dict = {
+ "user_id": _TEST_USER_ID,
+ "live_request": {
+ "input": "What is the exchange rate from USD to SEK?",
+ },
+ }
+
+ await request_queue.put(request_dict)
+ await request_queue.put(None) # Sentinel to end the stream.
+ events = []
+ async for event in app.bidi_stream_query(request_queue):
+ events.append(event)
+ assert len(events) == 1
+
  def test_create_session(self):
  app = reasoning_engines.AdkApp(
  agent=Agent(name=_TEST_AGENT_NAME, model=_TEST_MODEL)
@@ -568,3 +620,35 @@ async def test_async_stream_query_invalid_message_type(self):
  ):
  async for _ in app.async_stream_query(user_id=_TEST_USER_ID, message=123):
  pass
+
+ @pytest.mark.asyncio
+ async def test_bidi_stream_query_invalid_request_queue(self):
+ app = reasoning_engines.AdkApp(
+ agent=Agent(name=_TEST_AGENT_NAME, model=_TEST_MODEL)
+ )
+ request_queue = []
+ with pytest.raises(
+ TypeError,
+ match="request_queue must be an asyncio.Queue instance.",
+ ):
+ async for _ in app.bidi_stream_query(request_queue):
+ pass
+
+ @pytest.mark.asyncio
+ async def test_bidi_stream_query_invalid_first_request(self):
+ app = reasoning_engines.AdkApp(
+ agent=Agent(name=_TEST_AGENT_NAME, model=_TEST_MODEL)
+ )
+ request_queue = asyncio.Queue()
+ request_dict = {
+ "live_request": {
+ "input": "What is the exchange rate from USD to SEK?",
+ },
+ }
+ await request_queue.put(request_dict)
+ with pytest.raises(
+ ValueError,
+ match="The first request must have a user_id.",
+ ):
+ async for _ in app.bidi_stream_query(request_queue):
+ pass

tests/unit/vertex_langchain/test_agent_engines.py

@@ -108,6 +108,20 @@ def clone(self):
  return self
 
 
+class BidiStreamQueryEngine:
+ """A sample Agent Engine that implements `bidi_stream_query`."""
+
+ def set_up(self):
+ pass
+
+ async def bidi_stream_query(self, unused_request_queue) -> AsyncIterable[Any]:
+ """Runs the bidi stream engine."""
+ raise NotImplementedError()
+
+ def clone(self):
+ return self
+
+
 class OperationRegistrableEngine:
  """Add a test class that implements OperationRegistrable."""
 
@@ -141,6 +155,10 @@ async def async_stream_query(
  for chunk in _TEST_AGENT_ENGINE_STREAM_QUERY_RESPONSE:
  yield chunk
 
+ async def bidi_stream_query(self, unused_request_queue) -> AsyncIterable[Any]:
+ """Runs the bidi stream engine."""
+ raise NotImplementedError()
+
  # Add a custom method to test the custom stream method registration.
  def custom_stream_query(self, unused_arbitrary_string_name: str) -> Iterable[Any]:
  """Runs the stream engine."""
@@ -158,6 +176,12 @@ async def custom_async_stream_method(
  for chunk in _TEST_AGENT_ENGINE_STREAM_QUERY_RESPONSE:
  yield chunk
 
+ async def custom_bidi_stream_method(
+ self, unused_request_queue
+ ) -> AsyncIterable[Any]:
+ """Runs the bidi stream engine."""
+ raise NotImplementedError()
+
  def clone(self):
  return self
 
@@ -890,6 +914,17 @@ async def async_stream_query() -> str:
  return "RESPONSE"
 
 
+class InvalidCapitalizeEngineWithoutBidiStreamQuerySelf:
+ """A sample Agent Engine with an invalid bidi_stream_query method."""
+
+ def set_up(self):
+ pass
+
+ async def bidi_stream_query() -> AsyncIterable[Any]:
+ """Runs the engine."""
+ raise NotImplementedError()
+
+
 class InvalidCapitalizeEngineWithoutRegisterOperationsSelf:
  """A sample Agent Engine with an invalid register_operations method."""
 
@@ -1625,6 +1660,23 @@ def test_get_agent_framework(
  ),
  ),
  ),
+ (
+ "Update the bidi stream query engine",
+ {"agent_engine": BidiStreamQueryEngine()},
+ types.reasoning_engine_service.UpdateReasoningEngineRequest(
+ reasoning_engine=_generate_agent_engine_with_class_methods_and_agent_framework(
+ [],
+ _agent_engines._DEFAULT_AGENT_FRAMEWORK,
+ ),
+ update_mask=field_mask_pb2.FieldMask(
+ paths=[
+ "spec.package_spec.pickle_object_gcs_uri",
+ "spec.class_methods",
+ "spec.agent_framework",
+ ]
+ ),
+ ),
+ ),
  (
  "Update the operation registrable engine",
  {"agent_engine": OperationRegistrableEngine()},
@@ -2826,6 +2878,95 @@ async def test_async_stream_query_agent_engine_with_operation_schema(
  )
  )
 
+ # pytest does not allow absl.testing.parameterized.named_parameters.
+ @pytest.mark.parametrize(
+ "test_case_name, test_engine, test_class_method_docs, test_class_methods_spec",
+ [
+ (
+ "Default Bidi Stream Queryable (Not Operation Registrable) Engine",
+ BidiStreamQueryEngine(),
+ {},
+ _TEST_ASYNC_STREAM_QUERY_SCHEMAS,
+ ),
+ (
+ "Operation Registrable Engine",
+ OperationRegistrableEngine(),
+ {},
+ _TEST_OPERATION_REGISTRABLE_SCHEMAS,
+ ),
+ ],
+ )
+ @pytest.mark.asyncio
+ async def test_create_agent_engine_with_bidi_stream_query_operation_schema(
+ self,
+ test_case_name,
+ test_engine,
+ test_class_method_docs,
+ test_class_methods_spec,
+ ):
+ with mock.patch.object(
+ base.VertexAiResourceNoun,
+ "_get_gca_resource",
+ ) as get_gca_resource_mock:
+ test_spec = types.ReasoningEngineSpec()
+ test_spec.class_methods.extend(test_class_methods_spec)
+ get_gca_resource_mock.return_value = types.ReasoningEngine(
+ name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
+ spec=test_spec,
+ )
+ agent_engines.create(test_engine)
+
+ # pytest does not allow absl.testing.parameterized.named_parameters.
+ @pytest.mark.parametrize(
+ "test_case_name, test_engine, test_class_methods, test_class_methods_spec",
+ [
+ (
+ "Default Bidi Stream Queryable (Not Operation Registrable) Engine",
+ BidiStreamQueryEngine(),
+ [],
+ [],
+ ),
+ (
+ "Operation Registrable Engine",
+ OperationRegistrableEngine(),
+ [],
+ _TEST_OPERATION_REGISTRABLE_SCHEMAS,
+ ),
+ ],
+ )
+ @pytest.mark.asyncio
+ async def test_update_agent_engine_with_bidi_stream_query_operation_schema(
+ self,
+ test_case_name,
+ test_engine,
+ test_class_methods,
+ test_class_methods_spec,
+ update_agent_engine_mock,
+ ):
+ with mock.patch.object(
+ base.VertexAiResourceNoun,
+ "_get_gca_resource",
+ ) as get_gca_resource_mock:
+ test_spec = types.ReasoningEngineSpec()
+ test_spec.class_methods.append(_TEST_METHOD_TO_BE_UNREGISTERED_SCHEMA)
+ get_gca_resource_mock.return_value = types.ReasoningEngine(
+ name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=test_spec
+ )
+ test_agent_engine = agent_engines.create(MethodToBeUnregisteredEngine())
+ assert hasattr(test_agent_engine, _TEST_METHOD_TO_BE_UNREGISTERED_NAME)
+
+ with mock.patch.object(
+ base.VertexAiResourceNoun,
+ "_get_gca_resource",
+ ) as get_gca_resource_mock:
+ test_spec = types.ReasoningEngineSpec()
+ test_spec.class_methods.extend(test_class_methods_spec)
+ get_gca_resource_mock.return_value = types.ReasoningEngine(
+ name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
+ spec=test_spec,
+ )
+ test_agent_engine.update(agent_engine=test_engine)
+
 
 @pytest.mark.usefixtures("google_auth_mock")
 class TestAgentEngineErrors:
@@ -2887,8 +3028,8 @@ def test_create_agent_engine_no_query_method(
  TypeError,
  match=(
  "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`."
  ),
  ):
  agent_engines.create(
@@ -2911,8 +3052,8 @@ def test_create_agent_engine_noncallable_query_attribute(
  TypeError,
  match=(
  "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`."
  ),
  ):
  agent_engines.create(
@@ -3024,6 +3165,23 @@ def test_create_agent_engine_with_invalid_async_stream_query_method(
  requirements=_TEST_AGENT_ENGINE_REQUIREMENTS,
  )
 
+ def test_create_agent_engine_with_invalid_bidi_stream_query_method(
+ self,
+ create_agent_engine_mock,
+ cloud_storage_create_bucket_mock,
+ tarfile_open_mock,
+ cloudpickle_dump_mock,
+ cloudpickle_load_mock,
+ importlib_metadata_version_mock,
+ get_agent_engine_mock,
+ ):
+ with pytest.raises(ValueError, match="Invalid bidi_stream_query signature"):
+ agent_engines.create(
+ InvalidCapitalizeEngineWithoutBidiStreamQuerySelf(),
+ display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
+ requirements=_TEST_AGENT_ENGINE_REQUIREMENTS,
+ )
+
  def test_create_agent_engine_with_invalid_register_operations_method(
  self,
  create_agent_engine_mock,
@@ -3158,8 +3316,8 @@ def test_update_agent_engine_no_query_method(
  TypeError,
  match=(
  "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`."
  ),
  ):
  test_agent_engine = _generate_agent_engine_to_update()
@@ -3181,8 +3339,8 @@ def test_update_agent_engine_noncallable_query_attribute(
  TypeError,
  match=(
  "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`."
  ),
  ):
  test_agent_engine = _generate_agent_engine_to_update()
@@ -3324,7 +3482,8 @@ def test_update_class_methods_spec_with_registered_operation_not_found(self):
  "register the API methods: "
  "https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/develop/custom#custom-methods. "
  "Error: {Unsupported api mode: `UNKNOWN_API_MODE`, "
- "Supported modes are: ``, `async`, `async_stream`, `stream`.}"
+ "Supported modes are: ``, `async`, `async_stream`, "
+ "`bidi_stream`, `stream`.}"
  ),
  ),
  ],