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

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

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

https://pypi.org/project/litellm/

## 開発環境
 - Python 3.9
 - litellm 1.44.15

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

https://docs.litellm.ai/docs/providers/vertex

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

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

https://cloud.google.com/docs/authentication/provide-credentials-adc?hl=ja#google-idp

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

https://cloud.google.com/docs/authentication/application-default-credentials?hl=ja#personal

### Gemini 呼び出し
Google Cloud SDK では以下のように指定します。

``` python
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 で呼び出すには以下のように指定します。

``` python
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()` の引数に追加します。

``` python
    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'
```

## 参考

https://qiita.com/Yoshihiro-Hirose/items/9d12c3ffec5f9df15aae

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