🎯

LLM APIを良い感じに呼べればOKな時に便利なlitellm

2024/02/09に公開

こんにちは。ログラスのLLMチームでソフトウェアエンジニアをしているr-kagayaです。

LLMを使ったアプリケーション・機能を作りたいとなったらいくつかのライブラリ選択肢があります。代表例はLangChainでしょう。
最近はオブザーバリティツールとしてLangSmithが登場するなど、LLMシステムを構築する上で必要・便利なモジュールの網羅が進み、LLMアプリ開発のデファクトに向かって爆進しています。

プロダクションで使うことに批判的な声もありましたが、最近はver0.1(それでも0.1ですが)に到達。
LangChain、LangChain Community、LangChain Coreに切り出され、LangSmithやLCEL、LangGraphが登場したりと、LangChainとどう付き合うかの状況は変わっていきそうです。

https://analyticsindiamag.com/langchain-is-garbage-software/
https://tech-blog.abeja.asia/entry/advent-2023-day13
https://note.com/mahlab/n/n9d3742ae803a?magazine_key=m5c93525708cf

とはいえ、LangChainほど高機能・多機能なライブラリは不要で、良い感じに各LLMプロパイダーのAPIを呼べたらそれで十分なユースケースも存在することでしょう。
最近は上記のモチベーションの時に利用しているlitellmというライブラリの紹介をします。

https://litellm.ai/
https://github.com/BerriAI/litellm

litellmとは

litellmはOpenAI APIのフォーマットで、他プロパイダーのLLM APIを呼び出せるようにラップしてくれるオープンソースのライブラリです。

公式に記載されている、Call 100+ LLMs using the same Input/Output Formatが最も端的にlitellmのことを表しています。

litellmが使われているプロジェクトはいくつかあり、ChatGPTの「Advanced Data Analysis(旧Code Interpreter)」のOSS版実装であるOpenInterpreterやPRをAIレビューしてくれるCodium PR Agentもlitellmを利用してるようです。

https://docs.litellm.ai/docs/project

I/Fを気にせず各社のLLM APIを呼び出す

litellmの最も基本的な特徴です。
各社のLLM APIへのリクエストを同じフォーマットで叩けて、モニタリングも簡単に設定できます。

Completion APIへのリクエストとLangSmithへのリクエスト履歴の送信は以下記述だけで対応できます。

litellm.success_callback = ["langsmith"]
response = litellm.completion(
    model="gpt-3.5-turbo",
    messages=[
        {"role": "user", "content": "ハンバーグについて教えてください。"},
    ],
    metadata={
        "run_name": "litellm-completion-demo",
        "project_name": "litellm-completion-demo",
    }
)

これだけでLangSmithに連携されます。便利ですね。
(ただ直近はstreamモードだとLangSmithに連携できなくて、頭を悩ませています)

別モデルを利用したいときは、modelを差し替えるだけです。

response = litellm.completion(
    model="bedrock/anthropic.claude-instant-v1",
    messages=[
        {"role": "user", "content": "ハンバーグについて教えてください。"},
    ],
    metadata={
        "run_name": "litellm-completion-demo",
        "project_name": "litellm-completion-demo",
    }
)

現時点でlitellmが対応してるプロパイダーは以下です。
OpenAIからAzure、Huggingface、VertexAI、Gemini、Anthropic、AWS Bedrockと大抵の場合はサポートプロパイダーで困ることはないのではないでしょうか。

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

ローカルでLLMを動かすOllamaや検索エンジンにAIを組み込んだサービスであるPerplexity AIにも対応しています。
https://docs.litellm.ai/docs/providers/perplexity
https://docs.litellm.ai/docs/providers/ollama

全プロバイダー・モデルでOpenAI APIのフォーマットで扱えるだけでも結構便利です。
とはいえOllamaがOpenAI Chat Completion APIとの互換性があったりと、プロプライエタリモデル以外は、OpenAI APIのフォーマットになんだかんだ準拠することになる可能性はありそうだと思っています。
https://ollama.ai/blog/openai-compatibility

コールバックを用いてObservabilityツールと連携

liteLLMはコールバックとして、input_callbacks、success_callbacks、failure_callbacksを提供しています。

ロギングに関してもSentryやHelicornなど複数サービスをサポートしているため、コールバックに対応サービス名を渡すだけで基本的な連携は完了します。

litellm.success_callback = ["langsmith"]

https://docs.litellm.ai/docs/observability/langsmith_integration

カスタムコールバックも対応していて、モニタリングサービスとの連携以外でもカスタマイズも可能です。

https://docs.litellm.ai/docs/observability/custom_callback
https://docs.litellm.ai/docs/observability/slack_integration

複数LLMプロパイダー・モデルを用いたフォールバック

各社のLLM APIへのリクエストを同じフォーマットで叩けるので、基本的なフォールバック処理であれば簡潔に記述できるのも嬉しいです。

ChatGPTにお願いしたコードで恐縮ですが、モデル名を差し替えるだけで動かせるので、単純にモデルリストとループ機構を書けば簡易的なフォールバックを実現することができます。

def call_llm_with_fallback(user_message):
    model_fallback_list = ["gpt-4", "gemini/gemini-pro"]
    messages = [{"content": user_message, "role": "user"}]
    
    for model in model_fallback_list:
        try:
            response = litellm.completion(model=model, messages=messages)
            # 成功したらレスポンスを返す
            return response  
        except Exception as e:
            print(f"error occurred with {model}: {traceback.format_exc()}")
    
    # すべてのモデルで失敗した場合
    return "All model attempts failed."

user_message = "Hello, how are you?"
response = call_llm_with_fallback(user_message)
print(response)

ContextWindowを超過したことを表すExceptionも存在するので、各モデルでトークン数を超過した場合によりコンテキストウィンドウが大きいモデルにフォールバックする処理も簡単に書くことができます。

def completion_with_context_window_exceeded_fallback(user_message, initial_model, fallback_list):
    messages = [{"content": user_message, "role": "user"}]
    try:
        response = completion(model=initial_model, messages=messages)
        return response
    except ContextWindowExceededError:
        # 初期モデルでのコンテキストウィンドウの最大トークン数を取得
        initial_model_max_tokens = get_max_tokens(initial_model)
        for fallback_model in fallback_list:
            if initial_model_max_tokens < fallback_model["max_tokens"]:
                try:
                    response = completion(model=fallback_model["model"], messages=messages)
                    # フォールバックモデルで成功した場合、レスポンスを返す
                    return response  
                except ContextWindowExceededError:
                    # フォールバックモデルでも失敗した場合、次のモデルで再試行
                    continue
        # すべてのモデルで失敗した場合
        return "All model attempts failed due to context window exceeded."

context_window_fallback_list = [
    {"model": "gpt-3.5-turbo-16k", "max_tokens": 16385},
    {"model": "gpt-4-32k", "max_tokens": 32768},
    {"model": "claude-instant-1", "max_tokens": 100000}
]
user_message = "Hello, how are you?"
initial_model = "command-nightly"

# 関数を呼び出し
response = completion_with_context_window_exceeded_fallback(user_message, initial_model, context_window_fallback

カスタム価格の登録

上手く使えないか気になってるのが、モデルにカスタム価格を登録できる機能です。
トークンあたりのコスト、もしくは1秒あたりのコストに対して、カスタムでモデルアクセスの価格を登録できる機能のようです。
有料でLLMプロダクト・機能を用いてる場合に、上手く原価にアドオンした価格を設定して、コスト計算周りで使って省力化できればいいなーとか考えています。

下記のコード例だと、azureのモデルに対して、トークンあたりのコストを設定しています。

# azure call
response = completion(
  model = "azure/<your_deployment_name>", 
  messages = [{ "content": "Hello, how are you?","role": "user"}]
  input_cost_per_token=0.005,
  output_cost_per_token=1,
)
cost = completion_cost(completion_response=response)
print(cost)

https://docs.litellm.ai/docs/proxy/custom_pricing

LLM APIコールにバジェットを設定する

LLM APIコールに予算を設定できるBudget Managerという機能もあります。
カスタム価格と同じく、自前で作れと言われれば作りますが、よくあるユースケースとしてlitellm側で対応してくれるのは、地味に捗るいぶし銀な機能ではないかと考えています。
https://docs.litellm.ai/docs/budget_manager

ユーザー毎に利用可能な予算を設定して、それらを管理するためのクラスが提供されています。

budget_manager = litellm.BudgetManager(project_name="test_project-1")
user = "1"

# ユーザー毎に予算を設定する
if not budget_manager.is_valid_user(user):
    budget_manager.create_budget(total_budget=0.0001, user=user)

# 現在とトータルのbudgetを取得
current_cost = budget_manager.get_current_cost(user=user)
total_budget = budget_manager.get_total_budget(user)

# 1コール目: current_cost: 0.000000, total_budget: 0.000100
# 2コール目: current_cost: 0.000077, total_budget: 0.000100
# 3コール目: current_cost: 0.000152, total_budget: 0.000100 -> Sorry - no budget!
formatted_current_cost = f"{current_cost:.6f}"
formatted_total_budget = f"{total_budget:.6f}"
print(f"current_cost: {formatted_current_cost}")
print(f"total_budget: {formatted_total_budget}")

if current_cost <= total_budget:
    response = litellm.completion(model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hey, how's it going?"}])
    # APIコール完了後に利用分のコストを加算
    budget_manager.update_cost(completion_obj=response, user=user)
else:
    # ユーザーが予算を使い尽くした
    print("Sorry - no budget!")

実際に活用するためにはDB等に保存したり、リセットする処理が必要ですが、APIを通じて自社で抱えるDBに保存したり、

独自のDBを使用するには、BudgetManagerクライアントタイプをホスト型に設定し、api_baseを設定します。
あなたのAPIは/get_budgetと/set_budgetエンドポイントを公開することが期待されます。詳細はコードを見てください

クラウド版としてlitellm APIを用いることも出来るらしいです(こっちはOpenAI Proxy Serverを使う必要がありそう)

LiteLLM API は両方を提供します。ユーザーオブジェクトをホストされたデータベースに保存し、cron ジョブを毎日実行して、設定された期間に基づいてユーザー予算をリセットします (例: 予算を日次/週次/月次などにリセット)

https://docs.litellm.ai/docs/budget_manager#advanced-usage

おわりに

以上、litellmについて簡単に紹介しました。
今回は触れませんでしたが、キャッシュや、Load BalancingEmbeddingやImage GenerationのI/Fも提供されています。

LangChainに出来ないことはほぼないかもしれませんが、とはいえ良い感じにLLM APIを叩くことが出来たらOKな場合においては検討に値するライブラリではないかと思いました。

「LangChainはtoo muchだけど、OpenAIライブラリでAPI叩くだけなのも。。」という時にぜひlitellmのことを思い出してみてください!

株式会社ログラス テックブログ

Discussion