LiteLLMのカスタム・プロバイダ

概要

LiteLLM https://docs.litellm.ai/docs/ は、様々なLLM APIをラップして、共通の入力と出力を扱うことができるようにするためのPythonライブラリである。
ここではプロバイダの仕組みについて簡単にメモしておく。

基本的な使い方

たとえばchat completion (チャット補完) APIの呼び方はこう:

import litellm

response = litellm.completion(
            model="foo/bar/baz",
            api_base=...,
            api_key=...,
            api_version=...,  
            messages=[{"content": "Hello, how are you?", "role": "user"}],
            max_tokens=10
        )

print(response['choices'][0]['message']['content'])

モデル (model) とプロバイダ (provider)

モデル名の最初のセグメントをプロバイダと呼ぶ。APIベンダによって異なるプロバイダを指定する必要がある。
例えばmodel="foo/bar/baz"ならばプロバイダはfooである。
プロバイダ毎に、異なる実装により呼び出しが処理される。
その際に、モデル名は先頭のセグメントが取り除かれてbar/bazとして渡される。

openai_likeプロバイダ

最近の実装にはopenai_likeプロバイダが存在する。
これがOpenAIコンパチブルなAPIに使うものなのか、他のAPIに使うものなのか著者には不明である。

api_base

APIエンドポイントのURL。
ほとんどの場合、プロバイダに依存してデフォルト値が設定される。

例えばapi_base=https://foo/bar/bazならば、https://foo/bar/baz/chat/completionsのAPIが呼ばれる。

api_version

APIのバージョン。APIの入出力がバージョンにより異なる場合に使用すると思われる。

注意: ここで指定したバージョンはエンドポイントURLに影響を与えない。
つまり、1をこの属性値として設定しても、エンドポイントがhttps://<api_base>/v1などになるわけではない。
v1は自分でapi_baseに付加しておく必要がある。

api_key

APIの認証用のキー。例えばOpenAI APIではこの値がBearer認証に使われる。
この値はプロバイダ毎の環境変数でも設定できる。たとえばopenaiプロバイダであれば、OPENAI_API_KEYで設定できる。

実装の話をすると、これはグローバルなcompletion()メソッド内で保管されている。各プロバイダが処理しているわけではなさそう

カスタム・プロバイダ (詳細略)

特定のクラスを継承したクラスを定義し、それを登録することにより、カスタムのプロバイダが定義・使用できる。
Rerank APIでの例はこちら: https://docs.litellm.ai/docs/adding_provider/new_rerank_provider
Chat Completion APIでも同様に実装できると思われる。
ただし必要な処理を自前で実装し切るのは大変そう。