send_message_streaming_async fails with TypeError

[qxcv]

Environment details

• OS type and version: Ubuntu 22.04
• Python version: Python 3.10.11
• pip version: pip 23.3.1
• google-cloud-aiplatform version: tried at 760a025 and 8e42bc5

Steps to reproduce

See code example below.

Code example

Here is a simplified version of the code that is causing the issue:

async def call_palm_api(
    *, user: str, context: str, msg: str,
) -> AsyncGenerator[str, None]:
    global _VERTEX_INITIALIZED
    if not _VERTEX_INITIALIZED:
        vertexai.init(project="prompt-ad-game", location="us-central1")
        _VERTEX_INITIALIZED = True

    chat_model = ChatModel.from_pretrained("chat-bison@001")

    chat = chat_model.start_chat(
        context=context,
        max_output_tokens=max_tokens,
        temperature=0,
    )
    responses = chat.send_message_streaming_async(
        message=msg
    )

    async for resp in responses:
        yield resp.text

Stack trace

Traceback (most recent call last):
  File "/home/sam/anaconda3/envs/promptgame/lib/python3.10/site-packages/asgiref/sync.py", line 534, in thread_handler
    raise exc_info[1]
  File "/home/sam/anaconda3/envs/promptgame/lib/python3.10/site-packages/django/core/handlers/exception.py", line 42, in inner
    response = await get_response(request)
  File "/home/sam/anaconda3/envs/promptgame/lib/python3.10/site-packages/django/core/handlers/base.py", line 253, in _get_response_async
    response = await wrapped_callback(
  File "/home/sam/repos/promptgame/src/promptgame/gameui/views.py", line 419, in security_submit
    async for chunk, flag in promptgame.gameui.attack.streaming_llm_output(
  File "/home/sam/repos/promptgame/src/promptgame/gameui/attack.py", line 537, in streaming_llm_output
    next_text = await asyncio.wait_for(
  File "/home/sam/anaconda3/envs/promptgame/lib/python3.10/asyncio/tasks.py", line 445, in wait_for
    return fut.result()
  File "/home/sam/repos/promptgame/src/promptgame/gameui/llm_wrappers.py", line 233, in call_palm_api
    async for resp in responses:
  File "/home/sam/anaconda3/envs/promptgame/lib/python3.10/site-packages/vertexai/language_models/_language_models.py", line 2023, in send_message_streaming_async
    async for prediction_dict in _streaming_prediction.predict_stream_of_dicts_from_single_dict_async(
  File "/home/sam/anaconda3/envs/promptgame/lib/python3.10/site-packages/google/cloud/aiplatform/_streaming_prediction.py", line 242, in predict_stream_of_dicts_from_single_dict_async
    async for dict_list in predict_stream_of_dict_lists_from_single_dict_list_async(
  File "/home/sam/anaconda3/envs/promptgame/lib/python3.10/site-packages/google/cloud/aiplatform/_streaming_prediction.py", line 186, in predict_stream_of_dict_lists_from_single_dict_list_async
    async for tensor_list in predict_stream_of_tensor_lists_from_single_tensor_list_async(
  File "/home/sam/anaconda3/envs/promptgame/lib/python3.10/site-packages/google/cloud/aiplatform/_streaming_prediction.py", line 133, in predict_stream_of_tensor_lists_from_single_tensor_list_async
    async for response in prediction_service_async_client.server_streaming_predict(
TypeError: 'async for' requires an object with __aiter__ method, got coroutine

[qxcv]

I suspect the last call (in google/cloud/aiplatform/_streaming_prediction.py on line 133) just needs an await. i.e. convert

async for response in prediction_service_async_client.server_streaming_predict(

to

async for response in await prediction_service_async_client.server_streaming_predict(

Edit: yeah it's working for me now that I have the await.