generate_content_async broken in Colab?

[ochafik]

Description of the bug:

Calling GenerativeModel.generate_content_async in Colab results in the following error:

[/usr/local/lib/python3.10/dist-packages/google/generativeai/generative_models.py](https://localhost:8080/#) in generate_content_async(self, contents, generation_config, safety_settings, stream, tools, tool_config, request_options)
    383                 return await generation_types.AsyncGenerateContentResponse.from_aiterator(iterator)
    384             else:
--> 385                 response = await self._async_client.generate_content(
    386                     request,
    387                     **request_options,

[/usr/local/lib/python3.10/dist-packages/google/ai/generativelanguage_v1beta/services/generative_service/async_client.py](https://localhost:8080/#) in generate_content(self, request, model, contents, retry, timeout, metadata)
    404 
    405         # Send the request.
--> 406         response = await rpc(
    407             request,
    408             retry=retry,

TypeError: object GenerateContentResponse can't be used in 'await' expression

show colab cell code

!pip install -U -q google-generativeai
# installs google-generativeai==0.7.2

# The following two lines don't change anything:
# import nest_asyncio
# nest_asyncio.apply()

import asyncio
import google.generativeai as genai

from google.colab import userdata
GOOGLE_API_KEY=userdata.get('GOOGLE_API_KEY')

genai.configure(api_key=GOOGLE_API_KEY)

model = genai.GenerativeModel(model_name="models/gemini-1.5-pro-latest")

async def main():
  print(await model.generate_content_async("Tell me a joke"))

await main()

The ~same code works well outside Colab (tested Python 3.10.12 = same as Colab, 3.11 and 3.12)

show standalone code that works

'''
  This nearly identical code works w/ Python 3.10.12 (same as Colab), 3.11, 3.12

  conda create -n gemini-bug-3.10.12 -y
  conda activate gemini-bug-3.10.12
  pip install -U -q google-generativeai
  GOOGLE_API_KEY=... python bug.py
'''

import asyncio
import google.generativeai as genai

import os
GOOGLE_API_KEY=os.environ['GOOGLE_API_KEY']

genai.configure(api_key=GOOGLE_API_KEY)

model = genai.GenerativeModel(model_name="models/gemini-1.5-pro-latest")

async def main():
  print(await model.generate_content_async("Tell me a joke"))

if __name__ == "__main__":
  asyncio.run(main())

Full repro: https://gist.github.com/ochafik/f7472295d7b3f42824c8458f32bfa58b

Actual vs expected behavior:

Unexpected exception in Colab (works outside Colab)

Any other information you'd like to share?

No response

[MarkDaoust]

Interesting! Thanks for reporting this.

Colab sets genai.configure(transport="rest"), the problem seems to be async+rest transport.

In standard python this error is reproducible if I set transport="rest".

So that seems like an issue with the underlying generated client library.

[MarkDaoust]

This is: https://github.com/googleapis/gapic-generator-python/issues/1962

[ochafik]

FWIW, here's a workaround for the non-streamed case:

import aiohttp
import google.generativeai.protos as protos

class AsyncClient:
  def __init__(self, api_key):
    self._api_key = api_key

  async def generate_content(self, request, **request_options):
    async with aiohttp.ClientSession() as session:
      url = f"https://generativelanguage.googleapis.com/v1beta/{request.model}:generateContent?key={self._api_key}"
      async with session.post(url, data=protos.GenerateContentRequest.to_json(request)) as response:
        assert response.status == 200, f'Response status: {response.status}: {await response.text()}'
        return protos.GenerateContentResponse.from_json(await response.text())

model = genai.GenerativeModel(model_name="models/gemini-1.5-flash-latest")
model._async_client = AsyncClient(GOOGLE_API_KEY)

print(await model.generate_content_async("Hello, "))

[oryba]

There's another trick I've noticed in this official cookbok.

# This notebook should work fine in a normal Python environment, but due to https://github.com/google-gemini/generative-ai-python/issues/499
# this workaround is needed in Colab, effectively un-monkey-patching a Colab patch.
genai.configure = getattr(genai.configure, "func", genai.configure)

P.S. To be honest, working with Gemini ecosystem is a constant headache even when it comes to the basic things like this one. Literally everything from docs to official clients is extremely inconsistent.

[siriusmozzie48]

genai.configure = getattr(genai.configure, "func", genai.configure)

You are the best, THANK YOU !