Geminiが意図せず失敗したときの対処法

はじめに

Unipos Advent Calendar 2024、1日目の記事です！（late sub です！）

今回は Vertex AI の Gemini を使ったバッチ処理で実際に使ってる小ネタになります。Python SDK で Gemini などを触ってる人向けの記事になります。

そもそも Google Cloud 使ってるならもっと他にやり方はあるのですが、今回は一例として手軽にできる方法として紹介していきます。

背景

生成AIをfor文などで繰り返し使っていると、リクエスト数の上限や予期せぬエラーによって意図せず落ちることが多々ありました。

特に悩まされたエラーが

InternalServerError: 500 Internal error encountered.

です。皆さんも一度は経験したことがあるのではないでしょうか？

こいつに関して色々と調べていると、コミュニティにて同様の問題を提起してる記事に行き着きました。しかし今現在（2024/12/05）もまだ解決されてはいなさそうで、Googleのスタッフは「Pythonのバージョンを最新のものにアップグレードしてみたら？」とだけコメントしています。

なかなか言語そのもののバージョンアップは腰が重いので、他に方法はないのか〜と悩んでたところ、弊社の某スーパーエンジニアさんが「いい感じに再実行してくれるやつがあるんじゃない？」と鶴の一声をかけていただき、今回の対処法にたどり着きました。

対処法

それでは本題に入ります。

対処法は至ってシンプルで、以下のようにリクエストを飛ばしてる関数にデコレータをつけてあげるだけです。

import vertexai
from google.api_core import retry
from vertexai.generative_models import GenerativeModel, GenerationResponse

@retry.Retry()
def generate_content_with_retry(model: GenerativeModel, contents: str) -> GenerationResponse:
    return model.generate_content(contents)

if __name__ == "__main__":
    vertexai.init(project="my-project", location="my-location")
    model = GenerativeModel("gemini-pro")
    contents = "Why is sky blue?"
    response = generate_content_with_retry(model, contents)
    print(response)

また、リトライ時にログを飛ばしたい場合は以下のように設定することができます。

...
from logging import getLogger

def log_retry(retry_state: retry.retry_base._BaseRetry):
    getLogger().warning("Retry: %s", retry_state)

@retry.Retry(on_error=log_retry)
def generate_content_with_retry(model: GenerativeModel, contents: str) -> GenerationResponse:
    return model.generate_content(contents)
...

実際は上記の関数をfor文などで何回も実行する場合に役立つと思います。

おわりに

今回は Vertex AI の API 関連の小ネタを紹介しました。 Google Cloud では Vertex AI 含む、たくさんの便利な機能が登場してきているので、今後はそれらも触っていきながらアップデートしていければなと思います。

引き続きUnipos Advent Calendar 2024をお楽しみください。

参考

• Unipos Advent Calendar 2024
• Vertex Generative AI SDK for Python
• Retry