litellm を使って Vertex AI Gemini を呼び出す

はじめに

すでに GPT モデルを呼び出しているアプリケーションがあり、Vertex AI Gemini を追加することになりました。
リクエスト・レスポンスの構造が GPT モデルと異なるため、その差分を吸収するために litellm を使って開発しました。

開発環境

Python 3.9
litellm 1.44.15

開発手順

基本的に公式ドキュメントに沿って進めました。

開発環境準備

pip install google-cloud-aiplatform で必要なライブラリをインストールします。
認証設定は、サービスアカウントのキーをダウンロードして指定するか Google Cloud のアプリケーションデフォルト認証を使用するの2通りありますが、今回はアプリケーションデフォルト認証を使いました。

Python プロジェクトのディレクトリでコマンドプロンプトを起動し、gcloud auth application-default login コマンドを実行します。
ブラウザが自動的に開くので画面の案内通りに進めると認証ができます。

ちなみに gcloud auth application-default login コマンドで生成される認証ファイルは Windows の場合 %APPDATA%\gcloud\application_default_credentials.json に置かれるとのことです。

Gemini 呼び出し

Google Cloud SDK では以下のように指定します。

generation_config = {
    "max_output_tokens": 8192,
    "temperature": 0.4,
}

safety_settings = {
    generative_models.HarmCategory.HARM_CATEGORY_HATE_SPEECH: generative_models.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
    generative_models.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: generative_models.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
    generative_models.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: generative_models.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
    generative_models.HarmCategory.HARM_CATEGORY_HARASSMENT: generative_models.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
}

def generate_content():
    vertexai.init(project="<PROJECT ID>", location="us-central1")
    model = GenerativeModel("gemini-1.5-pro-001")
    responses = model.generate_content(
        ['What is Gemini?'],  # Gemini への質問をここで記述
        generation_config=generation_config,
        safety_settings=safety_settings,
        stream=True,
    )

    for response in responses:
        print(response.text, end="")

litellm で呼び出すには以下のように指定します。

def generate_content_litellm():
    os.environ["VERTEX_PROJECT"] = "<PROJECT ID>"
    os.environ["VERTEX_LOCATION"] = "us-central1"

    response = litellm.completion(
        model="vertex_ai/gemini-1.5-pro",
        messages=[
            {"role": "user", "content": "What is Gemini?"}
        ],
        stream=True,
        temperature=0.4,
        max_tokens=8192,
        safety_settings=[
            {
                "category": "HARM_CATEGORY_HARASSMENT",
                "threshold": "BLOCK_MEDIUM_AND_ABOVE",
            },
            {
                "category": "HARM_CATEGORY_HATE_SPEECH",
                "threshold": "BLOCK_MEDIUM_AND_ABOVE",
            },
            {
                "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
                "threshold": "BLOCK_MEDIUM_AND_ABOVE",
            },
            {
                "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
                "threshold": "BLOCK_MEDIUM_AND_ABOVE",
            },
        ],
    )
    
    for part in response:
        print(part.choices[0].delta.content or "")

completion() の引数 model の先頭に vertex_ai/ をつけることで、Vertex AI Gemini を使うことができます。

Google 検索によるグラウンディング

litellm で呼び出す場合でも、Google 検索を使ったグラウンディングを使うことができます。
tools に googleSearchRetrieval を追加し、 completion() の引数に追加します。

    tools = [{"googleSearchRetrieval": {"disable_attribution": False}}]

    response = litellm.completion(
        model="vertex_ai/gemini-1.5-pro",
        messages=[
            {"role": "user", "content": "Geminiについて教えて"}
        ],
        tools=tools,
        temperature=0.4,
        max_tokens=8192,
        safety_settings=[
            {
                "category": "HARM_CATEGORY_HARASSMENT",
                "threshold": "BLOCK_MEDIUM_AND_ABOVE",
            },
            {
                "category": "HARM_CATEGORY_HATE_SPEECH",
                "threshold": "BLOCK_MEDIUM_AND_ABOVE",
            },
            {
                "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
                "threshold": "BLOCK_MEDIUM_AND_ABOVE",
            },
            {
                "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
                "threshold": "BLOCK_MEDIUM_AND_ABOVE",
            },
        ],
    )
    print(response["choices"][0]["message"]["content"])

私が検証した時点では、Google 検索によるグラウンディングを使う場合、レスポンスのストリーミングはオフしておく必要がありました。
ストリーミングをオンにする（completion() の引数に stream=True を指定する）と、以下のエラーが発生しました。

Traceback (most recent call last):
  File "C:\<プロジェクトフォルダパス>\.venv\lib\site-packages\litellm\utils.py", line 10734, in __next__
    chunk = next(self.completion_stream)
  File "C:\<プロジェクトフォルダパス>\.venv\lib\site-packages\litellm\llms\vertex_ai_and_google_ai_studio\gemini\vertex_and_google_ai_studio_gemini.py", line 1576, in __next__
    return self.chunk_parser(chunk=json_chunk)
  File "C:\<プロジェクトフォルダパス>\.venv\lib\site-packages\litellm\llms\vertex_ai_and_google_ai_studio\gemini\vertex_and_google_ai_studio_gemini.py", line 1505, in chunk_parser
    gemini_chunk = processed_chunk["candidates"][0]
KeyError: 'candidates'
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
  File "C:\<プロジェクトフォルダパス>\.venv\lib\site-packages\litellm\utils.py", line 10822, in __next__
    raise exception_type(
  File "C:\<プロジェクトフォルダパス>\.venv\lib\site-packages\litellm\utils.py", line 8650, in exception_type
    raise e
  File "C:\<プロジェクトフォルダパス>\.venv\lib\site-packages\litellm\utils.py", line 8622, in exception_type
    raise APIConnectionError(
litellm.exceptions.APIConnectionError: litellm.APIConnectionError: 'candidates'
Traceback (most recent call last):
  File "C:\<プロジェクトフォルダパス>\.venv\lib\site-packages\litellm\utils.py", line 10734, in __next__
    chunk = next(self.completion_stream)
  File "C:\<プロジェクトフォルダパス>\.venv\lib\site-packages\litellm\llms\vertex_ai_and_google_ai_studio\gemini\vertex_and_google_ai_studio_gemini.py", line 1576, in __next__
    return self.chunk_parser(chunk=json_chunk)
  File "C:\<プロジェクトフォルダパス>\.venv\lib\site-packages\litellm\llms\vertex_ai_and_google_ai_studio\gemini\vertex_and_google_ai_studio_gemini.py", line 1505, in chunk_parser
    gemini_chunk = processed_chunk["candidates"][0]
KeyError: 'candidates'

参考

以上です。
どなたかの参考になれば幸いです。