はじめに

こんにちは。株式会社divxのエンジニア、渡部です。

今回は、私達が開発に携わっているAIチャットツールで遭遇した問題とその解決策についてご紹介しようと思います。

想像してみてください。正しく設定したはずのAPIキーが、知らない間に入れ替わっているとしたら…。

問題に遭遇するまでの経緯

何を開発していたのか

私のいるチームが開発しているのは、OpenAIのAPIを活用したAIチャットツールです。

単なるチャットだけでなく、AIコードレビューやブログ記事の校閲も拡張機能として備えています。

そして今回は、「社内独自情報をRAGで回答できる拡張機能」を開発していました。

通常のチャットと今回開発したRAG機能の使用状況を分析するため、それぞれで使うAPIキーを分けて実装しました。

参照させたい情報の保存と、その情報を元にした回答の生成に無事成功し、これで一件落着とチームメンバー全員が思っていました。

簡単に、RAGとは？

RAGとは、AIが学習していない情報を参照し、それを生成プロセスに統合する技術です。これにより、AIはより正確で関連性の高い回答を生成することが可能になります。

一般に公開されているAIは、世間に広く公開されていて、かつ、リアルタイム性の低い情報について回答するのは得意です。例えば、一般的な挨拶やプログラミング言語の仕様についてです。

しかし、ある特定の組織でだけ活用されていたり、最新の情報であったりすると、AIは学習していないため的外れな回答しかできません。

こういった情報もAIが答えられるよう、別のデータベースを用意してそこへAIに参照してほしい情報を保存します。

実際に回答を作るときは、AIに回答させる前に、まず保存された情報の中から回答に使えそうな情報を探し出します。

そして、使えそうな情報とユーザーからの質問をAIに渡すことで、学習していない情報についてもまるで知っていたかのように回答することができるのです。

事件発生

あと残すところはステージング環境での確認だけとなり、もう山場を超えたと思っていたところに、上司から突然連絡が入りました。

「通常チャットとRAGのAPIキー、入れ替わってない？」

私は耳を疑いました。そんなことがあるのかと。

しかし、APIの使用状況グラフを見るとたしかに入れ替わっているのです。

もう少し正確に表現すると、「圧倒的に使用回数の多いはずの通常チャットのAPIキー使用回数がほぼなくなり、新しく実装したRAGのAPIキー使用回数が跳ね上がっている」状況でした。

原因調査

入れ替わっている（？）事実は受け止めて調査開始です。

この時点では「誰か設定を間違えたのだろう」としか思っていませんでした。

しかし、調べていくとそんな単純な話ではありませんでした…。

疑わしき者たち

Parameter Store

APIキーの設定ミスといえば真っ先に浮かぶのが環境変数です。

AWSのSystems Manager Parameter Storeで環境変数を設定していましたので、その値を確認しにいきます。

ここで2つのAPIキーの値を間違って登録していれば入れ替え直せば済む話ですが、正しく設定されていました。本来なら喜ぶべきところなのですが、起きている現象の説明がつかないので混乱してしまいました。

GitLab

今回のアプリケーションのソースコードはGitLabで管理しており、CI/CDで使う環境変数の設定画面も触ることがあったので、次はそこを疑いました。

「何らかの原因でParameter Storeの値より優先されてしまったのでは？」という仮説のもと設定画面を確認しましたが、ここではそもそもAPIキーの設定をしていませんでした。

段々と雲行きが怪しくなってきました…。

Dockerfile & docker-compose.yml

環境変数の設定にミスはないとすると、次は環境変数の取得に問題がある可能性があります。

コンテナを使って開発しているためDockerfileまたはdocker-compose.ymlで見落としがあるかもしれません。

しかし、ここでもAPIキーの取得設定は正しかったのです。

config.py

コンテナでも環境変数の取得が間違っていないとなると、アプリケーションを疑わなければなりません。

PythonのフレームワークFlaskを使用しており、config.pyファイルで環境変数の値を定数として定義していました。

ここでの代入にミスがあると予想しましたが、結局は見当違いでした…。

幾多のメソッドたち

ここまでくると、基本的な設定は間違っておらず、実際にアプリケーションが動いている途中で入れ替わってしまっているとしか考えられません。

AIチャットで使っているソースコード全体に検索をかけて怪しい箇所を探しますが、間違えて違うAPIキーを渡していた場所は見つからず。APIキーは適切に管理されており、レビューで不具合を見落としていたわけでもなかったことが確認できたのは良かったですが、原因は相変わらず不明のままでした…。

原因特定

ここまで、普段の作業でよく編集するファイルを確認してきましたが原因は見つかりませんでした。

こうなると、我々が普段は編集しない場所でAPIキーの操作が行われている可能性があります。

我々以外でOpenAIのAPIキーを操作しているものがひとつだけあります。テキスト埋め込みで使用したライブラリchromadbです。

埋め込みに使うAPIキーとモデルの指定に以下のようなコードを使っていました。

from chromadb.utils import embedding_functions
from config import (
    EMBEDDING_MODEL,
    RAG_OPENAI_API_KEY,
)   
            
    embedding_function = embedding_functions.OpenAIEmbeddingFunction(
        api_key=RAG_OPENAI_API_KEY, model_name=EMBEDDING_MODEL
    )

chromaからembedding_functionsをインポートして、その中のOpenAIEmbeddingFunctionメソッドを使用していました。

このOpenAIEmbeddingFunctionメソッドでAPIキーはどのように扱われているのでしょうか？

実際のコードが以下になります。（GitHubはこちら）

import logging
from typing import Mapping, Optional, cast

from chromadb.api.types import Documents, EmbeddingFunction, Embeddings

logger = logging.getLogger(__name__)


class OpenAIEmbeddingFunction(EmbeddingFunction[Documents]):
    def __init__(
        self,
        api_key: Optional[str] = None,
        
        # 省略
        
        try:
            import openai
        except ImportError:
            raise ValueError(
                "The openai python package is not installed. Please install it with `pip install openai`"
            )
            
        if api_key is not None:
            openai.api_key = api_key
        # If the api key is still not set, raise an error
        elif openai.api_key is None:
            raise ValueError(
                "Please provide an OpenAI API key. You can get one at https://platform.openai.com/account/api-keys"
            )

api_keyが何かしら引数として渡された場合、それがopenai.api_keyとして設定されています。

ここで重要なのはインポートしたopenaiに直接api_keyを設定している部分です。

インスタンス化をせずに直接代入しているので、他のファイルで使用しているopenaiのapi_keyも上書きされてしまいます。

今回我々が遭遇したAPIキーが入れ替わっているように見える現象もこれが原因でした。

埋め込み機能を使用した瞬間にAPIキーが埋め込み用のものに上書きされてしまい、他の機能で使っているAPIキーが使えなくなっていたのです。

対策

アプリ内での対処

原因がわかったところで満足してはいけません。APIキーが上書きされてしまう問題を解決する必要があります。

上書きされないようにメソッドを自作するという案も出ましたが、これはかなりのコストがかかりそうです。ライブラリとの整合性をとりつつ自作していくのは大変です。

そこで、我々がとった策は「上書きされた直後に上書きし直そう」というシンプルな方法です。

まず、埋め込みを行うためのクラスを新たに作成しました。APIキーの再上書きを必ずセットにして行うためです。

import chromadb
import openai
from chromadb.utils import embedding_functions
from config import OPENAI_API_KEY

class SafeOpenAIEmbeddingFunction:

    def __init__(self, model_name: str, api_key: str):
        # OpenAIEmbeddingFunctionのインスタンスを作成
        self.embedding_function = embedding_functions.OpenAIEmbeddingFunction(
            api_key=api_key, model_name=model_name
        )
        # 元のAPIキーに上書きし直す
        openai.api_key = OPENAI_API_KEY

    def __call__(self, input):
        return self.embedding_function(input)

作成したこのクラスをもとに、一時的に使いたいRAG用のAPIキーを引数としてインスタンスを作成すれば問題は解決です。

embedding_function = SafeOpenAIEmbeddingFunction(
    api_key=OPENAI_API_KEY, model_name=EMBEDDING_MODEL
)

また、同様の問題が再発した場合に備えて、テキスト埋め込みの処理の直前と直後でopenai.api_keyの下4桁をログに出力させることにしました。こうすることで、APIキーが想定と異なる場合にはすぐ検知できるようになりました。

ライブラリでのIssue起票とPR提出

上書きされたから上書きし直すという、力技で無事に解決しました。

しかし、APIキーを使い分けたいと思う人は少ないかもしれませんが、一定の需要はあるはずです。ライブラリ側が修正されたらクラスを新規で作成する必要もないはずです。

ということで、新規Issueを立てたうえでPull Requestを作成してみました！

issue: https://github.com/chroma-core/chroma/issues/2979

PR: https://github.com/chroma-core/chroma/pull/3118

今回の修正では、インスタンス化したOpenAIEmbeddingFunctionのAPIキーに対して、引数として受け取ったキーの値を代入するようにしました。
引数で値を受け取れなかった場合はopenai.api_keyの値を流用します。

        self._api_key = api_key or openai.api_key
        # If the api key is still not set, raise an error
        if self._api_key is None:
            raise ValueError(

これにより、埋め込み直後にAPIキーを上書きする必要はなくなりました！

暫定対処で作成したクラスは不要になり、実際に埋め込みを実行するときは、ライブラリのメソッドを呼び出すだけでよくなりました。

from chromadb.utils import embedding_functions
from config import (
    EMBEDDING_MODEL,
    RAG_OPENAI_API_KEY,
)   

# 修正前
# from chroma_db_adapter import SafeOpenAIEmbeddingFunction ← 暫定対処で作成したクラスのインポート
#
#		 embedding_function = SafeOpenAIEmbeddingFunction(
#                api_key=RAG_CHATBOT_OPENAI_API_KEY, model_name=EMBEDDING_MODEL
#            )

#修正後      
    embedding_function = embedding_functions.OpenAIEmbeddingFunction(
        api_key=RAG_OPENAI_API_KEY, model_name=EMBEDDING_MODEL
    )

このPRは無事にマージされましたので、今後は同じ現象で悩む人はいなくなるはずです。

終わりに

今回の現象は「埋め込みとその他でAPIキーを使い分ける」という一般的ではないケースで発覚した、ライブラリ側の問題によるものでした。

複数のAPIキーを扱う際には、想定通りに使い分けができているか確認を徹底する必要があります。また、ライブラリのコードを読んでおくことの大切さも改めて実感しました。

今回の調査を通してOSSライブラリの改善点を見つけられたのは、良い経験になったと思います。

我々DIVXは、このような修正を通してOSSへの貢献も続けていきます。