Skip to content

Commit 99d76d1

Browse files
author
uncured7036
committed
fix: make async_stream_query fully async
1 parent 74e52d9 commit 99d76d1

File tree

2 files changed

+4
-4
lines changed

2 files changed

+4
-4
lines changed

google/cloud/aiplatform_v1/services/reasoning_engine_execution_service/async_client.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -410,7 +410,7 @@ async def sample_query_reasoning_engine():
410410
# Done; return the response.
411411
return response
412412

413-
def stream_query_reasoning_engine(
413+
async def stream_query_reasoning_engine(
414414
self,
415415
request: Optional[
416416
Union[
@@ -543,7 +543,7 @@ async def sample_stream_query_reasoning_engine():
543543
self._client._validate_universe_domain()
544544

545545
# Send the request.
546-
response = rpc(
546+
response = await rpc(
547547
request,
548548
retry=retry,
549549
timeout=timeout,

vertexai/agent_engines/_agent_engines.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1533,14 +1533,14 @@ def _wrap_async_stream_query_operation(
15331533
"""
15341534

15351535
async def _method(self, **kwargs) -> AsyncIterable[Any]:
1536-
response = self.execution_api_client.stream_query_reasoning_engine(
1536+
response = await self.execution_async_client.stream_query_reasoning_engine(
15371537
request=aip_types.StreamQueryReasoningEngineRequest(
15381538
name=self.resource_name,
15391539
input=kwargs,
15401540
class_method=method_name,
15411541
),
15421542
)
1543-
for chunk in response:
1543+
async for chunk in response:
15441544
for parsed_json in _utils.yield_parsed_json(chunk):
15451545
if parsed_json is not None:
15461546
yield parsed_json

0 commit comments

Comments
 (0)