Skip to content

Commit d62114d

Browse files
vertex-sdk-botcopybara-github
authored andcommitted
fix: make async_stream_query non-blocking by using the async execution client
PiperOrigin-RevId: 987300692
1 parent 5b0e8bb commit d62114d

2 files changed

Lines changed: 118 additions & 8 deletions

File tree

‎tests/unit/vertex_langchain/test_agent_engines.py‎

Lines changed: 112 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414
#
15+
import asyncio
1516
import cloudpickle
1617
import difflib
1718
import importlib
@@ -20,6 +21,7 @@
2021
import sys
2122
import tarfile
2223
import tempfile
24+
import time
2325
from typing import Any, AsyncIterable, Dict, Iterable, List, Optional
2426
from unittest import mock
2527

@@ -847,6 +849,37 @@ def mock_streamer():
847849
yield stream_query_agent_engine_mock
848850

849851

852+
class _AsyncChunkIterator:
853+
"""Yields response chunks asynchronously to simulate grpc.aio streams."""
854+
855+
def __init__(self, chunks):
856+
self._chunks = iter(chunks)
857+
858+
def __aiter__(self):
859+
return self
860+
861+
async def __anext__(self):
862+
try:
863+
return next(self._chunks)
864+
except StopIteration:
865+
raise StopAsyncIteration
866+
867+
868+
@pytest.fixture(scope="function")
869+
def async_stream_query_agent_engine_mock():
870+
# Simulates the GAPIC async client contract of returning an awaitable that
871+
# resolves to an async iterator.
872+
async def mock_streamer(*args, **kwargs):
873+
return _AsyncChunkIterator(_TEST_AGENT_ENGINE_STREAM_QUERY_RESPONSE)
874+
875+
with mock.patch.object(
876+
reasoning_engine_execution_service.ReasoningEngineExecutionServiceAsyncClient,
877+
"stream_query_reasoning_engine",
878+
side_effect=mock_streamer,
879+
) as async_stream_query_agent_engine_mock:
880+
yield async_stream_query_agent_engine_mock
881+
882+
850883
@pytest.fixture(scope="function")
851884
def get_gca_resource_mock():
852885
with mock.patch.object(
@@ -2835,7 +2868,7 @@ async def test_async_stream_query_after_create_agent_engine_with_operation_schem
28352868
test_engine,
28362869
test_class_method_docs,
28372870
test_class_methods_spec,
2838-
stream_query_agent_engine_mock,
2871+
async_stream_query_agent_engine_mock,
28392872
):
28402873
with mock.patch.object(
28412874
base.VertexAiResourceNoun,
@@ -2857,7 +2890,7 @@ async def test_async_stream_query_after_create_agent_engine_with_operation_schem
28572890

28582891
assert len(results) == 2 # Matches the length of mocked response
28592892

2860-
stream_query_agent_engine_mock.assert_called_with(
2893+
async_stream_query_agent_engine_mock.assert_called_with(
28612894
request=types.StreamQueryReasoningEngineRequest(
28622895
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
28632896
input={"input": _TEST_QUERY_PROMPT},
@@ -2894,7 +2927,7 @@ async def test_async_stream_query_after_update_agent_engine_with_operation_schem
28942927
test_class_methods,
28952928
test_class_methods_spec,
28962929
update_agent_engine_mock,
2897-
stream_query_agent_engine_mock,
2930+
async_stream_query_agent_engine_mock,
28982931
):
28992932
with mock.patch.object(
29002933
base.VertexAiResourceNoun,
@@ -2908,6 +2941,7 @@ async def test_async_stream_query_after_update_agent_engine_with_operation_schem
29082941
test_agent_engine = agent_engines.create(MethodToBeUnregisteredEngine())
29092942
assert hasattr(test_agent_engine, _TEST_METHOD_TO_BE_UNREGISTERED_NAME)
29102943

2944+
async_client_before_update = test_agent_engine.execution_async_client
29112945
with mock.patch.object(
29122946
base.VertexAiResourceNoun,
29132947
"_get_gca_resource",
@@ -2920,6 +2954,10 @@ async def test_async_stream_query_after_update_agent_engine_with_operation_schem
29202954
)
29212955
test_agent_engine.update(agent_engine=test_engine)
29222956

2957+
# Ensures update() reinitializes the async client used for streaming queries.
2958+
assert (
2959+
test_agent_engine.execution_async_client is not async_client_before_update
2960+
)
29232961
assert not hasattr(test_agent_engine, _TEST_METHOD_TO_BE_UNREGISTERED_NAME)
29242962
for method_name in test_class_methods:
29252963
invoked_method = getattr(test_agent_engine, method_name)
@@ -2929,7 +2967,7 @@ async def test_async_stream_query_after_update_agent_engine_with_operation_schem
29292967

29302968
assert len(results) == 2 # Matches the length of mocked response
29312969

2932-
stream_query_agent_engine_mock.assert_called_with(
2970+
async_stream_query_agent_engine_mock.assert_called_with(
29332971
request=types.StreamQueryReasoningEngineRequest(
29342972
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
29352973
input={"input": _TEST_QUERY_PROMPT},
@@ -2965,7 +3003,7 @@ async def test_async_stream_query_agent_engine_with_operation_schema(
29653003
test_engine,
29663004
test_class_methods,
29673005
test_class_methods_spec,
2968-
stream_query_agent_engine_mock,
3006+
async_stream_query_agent_engine_mock,
29693007
):
29703008
with mock.patch.object(
29713009
base.VertexAiResourceNoun,
@@ -2987,14 +3025,82 @@ async def test_async_stream_query_agent_engine_with_operation_schema(
29873025

29883026
assert len(results) == 2 # Matches the length of mocked response
29893027

2990-
stream_query_agent_engine_mock.assert_called_with(
3028+
async_stream_query_agent_engine_mock.assert_called_with(
29913029
request=types.StreamQueryReasoningEngineRequest(
29923030
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
29933031
input={"input": _TEST_QUERY_PROMPT},
29943032
class_method=method_name,
29953033
)
29963034
)
29973035

3036+
@pytest.mark.asyncio
3037+
async def test_async_stream_query_keeps_event_loop_responsive(self):
3038+
num_chunks = 5
3039+
chunk_delay_s = 0.02
3040+
watchdog_tick_s = 0.001
3041+
3042+
class _SlowAsyncStream:
3043+
"""Yields chunks with non-blocking delays to simulate a grpc.aio stream."""
3044+
3045+
def __init__(self):
3046+
self._remaining = num_chunks
3047+
3048+
def __aiter__(self):
3049+
return self
3050+
3051+
async def __anext__(self):
3052+
if self._remaining == 0:
3053+
raise StopAsyncIteration
3054+
self._remaining -= 1
3055+
await asyncio.sleep(chunk_delay_s)
3056+
return _TEST_AGENT_ENGINE_STREAM_QUERY_RESPONSE[0]
3057+
3058+
async def mock_async_streamer(*args, **kwargs):
3059+
return _SlowAsyncStream()
3060+
3061+
def mock_blocking_sync_streamer(*args, **kwargs):
3062+
"""Simulates a blocking sync client that blocks the thread during stream iteration."""
3063+
for _ in range(num_chunks):
3064+
time.sleep(chunk_delay_s)
3065+
yield _TEST_AGENT_ENGINE_STREAM_QUERY_RESPONSE[0]
3066+
3067+
test_agent_engine = mock.MagicMock()
3068+
test_agent_engine.resource_name = _TEST_AGENT_ENGINE_RESOURCE_NAME
3069+
test_agent_engine.execution_async_client.stream_query_reasoning_engine = (
3070+
mock_async_streamer
3071+
)
3072+
test_agent_engine.execution_api_client.stream_query_reasoning_engine = (
3073+
mock_blocking_sync_streamer
3074+
)
3075+
3076+
ticks = 0
3077+
stream_finished = asyncio.Event()
3078+
3079+
async def watchdog():
3080+
nonlocal ticks
3081+
while not stream_finished.is_set():
3082+
await asyncio.sleep(watchdog_tick_s)
3083+
ticks += 1
3084+
3085+
watchdog_task = asyncio.create_task(watchdog())
3086+
invoked_method = _agent_engines._wrap_async_stream_query_operation(
3087+
method_name=_TEST_DEFAULT_ASYNC_STREAM_METHOD_NAME
3088+
)
3089+
results = [
3090+
chunk
3091+
async for chunk in invoked_method(
3092+
test_agent_engine, input=_TEST_QUERY_PROMPT
3093+
)
3094+
]
3095+
stream_finished.set()
3096+
await watchdog_task
3097+
3098+
assert len(results) == num_chunks
3099+
# A blocking client starves the watchdog task by never yielding to the
3100+
# event loop. The threshold uses a loose lower bound to tolerate CI
3101+
# scheduling delays while verifying loop responsiveness.
3102+
assert ticks > num_chunks * 2
3103+
29983104
# pytest does not allow absl.testing.parameterized.named_parameters.
29993105
@pytest.mark.parametrize(
30003106
"test_case_name, test_engine, test_class_method_docs, test_class_methods_spec",

‎vertexai/agent_engines/_agent_engines.py‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -892,6 +892,9 @@ def update(
892892
self.execution_api_client = initializer.global_config.create_client(
893893
client_class=aip_utils.AgentEngineExecutionClientWithOverride,
894894
)
895+
self.execution_async_client = initializer.global_config.create_client(
896+
client_class=aip_utils.AgentEngineExecutionAsyncClientWithOverride,
897+
)
895898
# We use `._get_gca_resource(...)` instead of `created_resource` to
896899
# fully instantiate the attributes of the agent engine.
897900
self._gca_resource = self._get_gca_resource(resource_name=self.resource_name)
@@ -1730,14 +1733,15 @@ def _wrap_async_stream_query_operation(
17301733
"""
17311734

17321735
async def _method(self, **kwargs) -> AsyncIterable[Any]:
1733-
response = self.execution_api_client.stream_query_reasoning_engine(
1736+
response = await self.execution_async_client.stream_query_reasoning_engine(
17341737
request=aip_types.StreamQueryReasoningEngineRequest(
17351738
name=self.resource_name,
17361739
input=kwargs,
17371740
class_method=method_name,
17381741
),
17391742
)
1740-
for chunk in response:
1743+
async for chunk in response:
1744+
# In-memory chunk parsing requires no I/O.
17411745
for parsed_json in _utils.yield_parsed_json(chunk):
17421746
if parsed_json is not None:
17431747
yield parsed_json

0 commit comments

Comments
 (0)