LLM APIを良い感じに呼べればOKな時に便利なlitellm
こんにちは。ログラスのLLMチームでソフトウェアエンジニアをしているr-kagayaです。
LLMを使ったアプリケーション・機能を作りたいとなったらいくつかのライブラリ選択肢があります。代表例はLangChainでしょう。
最近はオブザーバリティツールとしてLangSmithが登場するなど、LLMシステムを構築する上で必要・便利なモジュールの網羅が進み、LLMアプリ開発のデファクトに向かって爆進しています。
プロダクションで使うことに批判的な声もありましたが、最近はver0.1(それでも0.1ですが)に到達。
LangChain、LangChain Community、LangChain Coreに切り出され、LangSmithやLCEL、LangGraphが登場したりと、LangChainとどう付き合うかの状況は変わっていきそうです。
とはいえ、LangChainほど高機能・多機能なライブラリは不要で、良い感じに各LLMプロパイダーのAPIを呼べたらそれで十分なユースケースも存在することでしょう。
最近は上記のモチベーションの時に利用している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を利用してるようです。
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と大抵の場合はサポートプロパイダーで困ることはないのではないでしょうか。
ローカルでLLMを動かすOllamaや検索エンジンにAIを組み込んだサービスであるPerplexity AIにも対応しています。
全プロバイダー・モデルでOpenAI APIのフォーマットで扱えるだけでも結構便利です。
とはいえOllamaがOpenAI Chat Completion APIとの互換性があったりと、プロプライエタリモデル以外は、OpenAI APIのフォーマットになんだかんだ準拠することになる可能性はありそうだと思っています。
コールバックを用いてObservabilityツールと連携
liteLLMはコールバックとして、input_callbacks、success_callbacks、failure_callbacksを提供しています。
ロギングに関してもSentryやHelicornなど複数サービスをサポートしているため、コールバックに対応サービス名を渡すだけで基本的な連携は完了します。
litellm.success_callback = ["langsmith"]
カスタムコールバックも対応していて、モニタリングサービスとの連携以外でもカスタマイズも可能です。
複数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)
LLM APIコールにバジェットを設定する
LLM APIコールに予算を設定できるBudget Managerという機能もあります。
カスタム価格と同じく、自前で作れと言われれば作りますが、よくあるユースケースとしてlitellm側で対応してくれるのは、地味に捗るいぶし銀な機能ではないかと考えています。
ユーザー毎に利用可能な予算を設定して、それらを管理するためのクラスが提供されています。
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 ジョブを毎日実行して、設定された期間に基づいてユーザー予算をリセットします (例: 予算を日次/週次/月次などにリセット)
おわりに
以上、litellmについて簡単に紹介しました。
今回は触れませんでしたが、キャッシュや、Load Balancing、EmbeddingやImage GenerationのI/Fも提供されています。
LangChainに出来ないことはほぼないかもしれませんが、とはいえ良い感じにLLM APIを叩くことが出来たらOKな場合においては検討に値するライブラリではないかと思いました。
「LangChainはtoo muchだけど、OpenAIライブラリでAPI叩くだけなのも。。」という時にぜひlitellmのことを思い出してみてください!
Discussion