Zenn
🐶

Difyのコードを読み解く！LLMのAPIキーはどのように保存されているのか？

2024/09/05に公開
AI
Python
OpenAI
Dify
tech

はじめに

最近Difyの勉強会も増え、ソースコードを読んでいましたが、気になるところが多すぎる!!! ということで、1つずつ記事にまとめて疑問を潰していこうと思っています。

今回はOpenAIやAnthropicのAPIキーがどのようにどこに保存されているのか、ソースコードを読み解いたのでまとめてみました。

これ知りたい!!!のリクエストお待ちしています

  • この実装気になってるけどどうなってる？
  • Difyのここって改善できる？

などありましたら気軽にコメントください。ソッコーで調べて記事にします...（たぶん
ということで今回は「LLMのAPIキーはどのように保存されているか」を見てみました。

version

  • Dify: v0.7.2

まずはUIを確認

OpenAIのAPIキーを発行し、Difyに登録します。

UIの操作としては、APIキーを入力し、保存ボタンを押して完了です。

DBを確認

DBに接続し、どのように保存されているかを確認してみます。

DifyのDBに接続する

Docker DesktopのUIからPostgresに接続してみます。
※環境変数が.env.exampleから変更していない場合は以下のコマンドで接続できると思います。変更している場合は、自分が設定した値に変更して実行してください。

$ psql -U postgres -d dify

テーブルに保存されている値を確認する

Difyには現在71テーブル存在していますが、各プロバイダーの情報はprovidersテーブルに保存されております。テーブルの情報は以下のとおりです。

dify=# \d providers
                                      Table "public.providers"
      Column      |            Type             | Collation | Nullable |           Default           
------------------+-----------------------------+-----------+----------+-----------------------------
 id               | uuid                        |           | not null | uuid_generate_v4()
 tenant_id        | uuid                        |           | not null | 
 provider_name    | character varying(255)      |           | not null | 
 provider_type    | character varying(40)       |           | not null | 'custom'::character varying
 encrypted_config | text                        |           |          | 
 is_valid         | boolean                     |           | not null | false
 last_used        | timestamp without time zone |           |          | 
 quota_type       | character varying(40)       |           |          | ''::character varying
 quota_limit      | bigint                      |           |          | 
 quota_used       | bigint                      |           |          | 
 created_at       | timestamp without time zone |           | not null | CURRENT_TIMESTAMP(0)
 updated_at       | timestamp without time zone |           | not null | CURRENT_TIMESTAMP(0)
Indexes:
    "provider_pkey" PRIMARY KEY, btree (id)
    "provider_tenant_id_provider_idx" btree (tenant_id, provider_name)
    "unique_provider_name_type_quota" UNIQUE CONSTRAINT, btree (tenant_id, provider_name, provider_type, quota_type)

APIキーのクレデンシャルはencrypted_configのカラムに入っています。
※こちらのAPIキーはすでに無効化している&途中の文字を変えています。

dify=# SELECT * FROM providers;
                  id                  |              tenant_id               | provider_name | provider_type |                                                                                                                                                                                                                                                                                        encrypted_config                                                                                                                                                                                                                                                                                        | is_valid |         last_used          | quota_type | quota_limit | quota_used |     created_at      |     updated_at      
--------------------------------------+--------------------------------------+---------------+---------------+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------+----------------------------+------------+-------------+------------+---------------------+---------------------
 f4af37da-e694-4261-a04c-a6492ab463e0 | b39cb7a5-99a9-4011-bfd6-ee070efcffcb | anthropic     | custom        | {"anthropic_api_key": "SFlCUklEOg4P9RndXYg4zyr3iLppiSvOzDOmozkwX0iz4BmMP+tAO8T+RefjKit5HciYPTHI8NN909dGiktGPACGTpHdZxCuFz7x0omm1vaTWYFYajlaF+ZAD+KWJeZ8BaAO13ujL5ZNDryWzsYAvPAf5h11zTY7bkn6OoEnTs0NDjhIFr5D1KRph2zhvOXWHxiMLPSZxue55r2t11oDP5n7Ez8kaZ1Oui2c5Ls1/rQXe9OJmO0/BHufF4zb1jdX5oOmgRw+kvEF6gN9Q6dY1ZckB4QXfOv+SCIhBCTe7cWiAezRcEyaAbzPSajIEjz2Xcs7Je4zF3b+S3Yrr4//JgD9ApJZcKD0W8WbI2It1Qqxk8w+iPNYUaAmNCM1siMcOKwJf5Ht31922umDAbigNPFtG+6BH2El2ESgTG716s1XuLZhZNnWOwcn6JtWu/1GZm0nGRhq1w6T3ywCn3o+iqLtfCBjmzsc3PoU5DROjX3zk6aMruNqOUGpIUN55DYpduV9kE5MomWgtjsx8ZJ/ZT8LIg=="}          | t        |                            |            |             |          0 | 2024-09-02 15:43:07 | 2024-09-02 15:43:07
 f350eb6a-04f3-4f2f-96e9-0e395c1735ad | b39cb7a5-99a9-4011-bfd6-ee070efcffcb | openai        | custom        | {"openai_api_key": "SFlCUklEOoqn1FHuMZbVqNNZlyB6xry4KVKOE0Y4Z9brHQfydcpSAhK8hsav8KNqLiMsU3/yDI7fsXgVvwvYO9ioNC939XkaokgtpENQSun1dOYGtQrVBE39jHNcUVn+5U1Dnu1kSgWdED7xpo75GMvbjH/UDx8A+ADqvH+4g9c0ylyzozi468VZY2BWIruKSegl757KvowRwMAFwL4t8oIeQvCGeFCqVz//o5csRoukVTwgm296hmz0PtRh1ee5wsliIS8Ozng3L59DCM8wUGHo2zUrWPblmZ8A4KAf8OKLfMMf5S2hpGpXJwAMXgTBH+OvVZ3HD5pAkCDMWtBV3S6o5wIKjzZ20dW89jSFWMVs6+U0+pz4+zgIZ4TZuThh22BfBOtj1nlyEq2964fGGCbN9nCQ1irRbSvO8eqga07FYmsmv2R/sZkEH+VcdlpDut7hCTiJTX7DT2aEMgHYu5i7cQC0XQJgyQtOxilaDmFBdnV2yFhO4ZYHsKOLN4sCicToPSXmN3XphRShdy/1v3E1vW5LRPHcWHB0mw=="} | t        | 2024-09-03 03:41:04.612645 |            |             |          0 | 2024-09-02 15:35:42 | 2024-09-02 15:35:42

今回はOpenAI, Anthropicの2つのキーを設定したためレコードが2つできており、以下のようなJSONフォーマットのテキストで格納されていることが分かります。

{"openai_api_key": "xxxxxxxxxx"}

暗号化処理の実装を見てみる

暗号化の処理は以下で実装されているため、中身を詳しく見ていきます。
このメソッド内で使っている関数などは、この後色々出てきます。

https://github.com/langgenius/dify/blob/dadca0f91a790572e06dbc127349c73cd6eec62f/api/core/entities/provider_configuration.py#L167-L218

このメソッドでは、主に以下の2つの構成となっています。

  1. 既存の認証情報の復号化
  2. 新しい認証情報の暗号化

順を追って説明していきます。

0. 準備: プロバイダー情報の取得

本題に入る前に、まずは暗号化処理に必要な準備段階のコードを見ていきます。
データベースからプロバイダー（Provider）の情報を取得し、以下3つの情報でフィルタリングしています。

  • tenant_id
  • provider_name
  • provider_type
        # get provider
        provider_record = db.session.query(Provider) \
            .filter(
            Provider.tenant_id == self.tenant_id,
            Provider.provider_name == self.provider.provider,
            Provider.provider_type == ProviderType.CUSTOM.value
        ).first()

次の処理は、プロバイダーの認証情報スキーマから"秘密の"変数を抽出するためのものです。

provider_credential_secret_variables = self.extract_secret_variables(
    self.provider.provider_credential_schema.credential_form_schemas
    if self.provider.provider_credential_schema else []
)

実際には以下のような値が入っています。

provider_credential_secret_variables=['openai_api_key']

1. 既存の認証情報の復号化

0. 準備: プロバイダー情報の取得でプロバイダーとプロバイダー固有の変数が取得できたので、復号化処理を見ていきます。

if provider_record.encrypted_config:
    if not provider_record.encrypted_config.startswith("{"):
        original_credentials = {
            "openai_api_key": provider_record.encrypted_config
        }
    else:
        original_credentials = json.loads(provider_record.encrypted_config)
else:
    original_credentials = {}

ここでは、provider_recordに保存されている暗号化された設定を取得しています。もし暗号化された設定が存在し、JSONフォーマットでない場合（"{"で始まらない場合）、それをOpenAIのAPIキーとして扱います。JSONフォーマットの場合は、それをデコードします。

すでにAPIキーが登録されている場合、provider_record.encrypted_configには以下のような値が入ります。この値に対して"{"で開始するかどうかの条件分岐が入ります。

provider_record.encrypted_config='{"openai_api_key": "SFlCUklEOgX0JiAKw...xxx"}'

その後、以下のコードで秘密の認証情報を復号化します。

# if send [__HIDDEN__] in secret input, it will be same as original value
if value == HIDDEN_VALUE and key in original_credentials:
    credentials[key] = encrypter.decrypt_token(self.tenant_id, original_credentials[key])

HIDDEN_VALUE は特別な定数値で、"[__HIDDEN__]"という文字列です。これは、ユーザーインターフェイスでパスワードなどの機密情報を "****" のように表示する際に使用されます。

HIDDENが入ってきた場合(値がAPIキーの場合) & keyがoriginal_credentials(openai_api_keyなど)の中にある場合は、復号化します。

なぜ復号化する？

後からOrganizationなどを追加したときに、既存のAPIキーを再度入力しなくてもいいように、このような実装になっていると推測しています。一度復号化して平文に戻し、そこに新しい情報を追加して再度暗号化するためではないかと。

※ちなみに、秘密鍵はRedisの中に入っていそうです。

https://github.com/langgenius/dify/blob/3230f4a0ec295095d2a25abeebcb4f3b1c84bf71/api/libs/rsa.py#L29-L90

2. 新しい認証情報の暗号化

for key, value in credentials.items():
    if key in provider_credential_secret_variables:
        credentials[key] = encrypter.encrypt_token(self.tenant_id, value)

ここでは、provider_credential_secret_variablesに含まれるキーに対応する値をencrypter.encrypt_token()メソッドを使用して暗号化しています。

https://github.com/langgenius/dify/blob/3230f4a0ec295095d2a25abeebcb4f3b1c84bf71/api/core/helper/encrypter.py#L15-L25

RSA暗号化アルゴリズムを使用してトークンを暗号化します。また、テナント固有の公開鍵（encrypt_public_key）を使用し各テナントのデータが個別に保護されます。

self.tenant_idが暗号化/復号化の際に使用されていることから、テナントごとに異なる暗号化が行われていることがわかります。

まとめ

  • LLMのクレデンシャルはprovidersテーブルのencrypted_configカラムに{"openai_api_key": "xxxxxxxxxx"}のような形式で暗号化されて保存されている
  • Providerごと、tenant_idごとに保存されている

以上、最後まで読んでいただきありがとうございます🙇‍♂️

Discussion