RAG(Retrieval Augmented Generation)で、ドキュメントの検索対象をユーザー権限に応じて制限して回答生成させたい場合の仕組みについて説明します。
概要
Azure Data Lake Storage Gen2(以下ADLS)にはACLによるアクセス制御をすることができます。
このACLを利用して、各ファイルに付与されているACLのoidをもとにAzure AI Search(以下AI Search)の検索フィルターにユーザーのoidが含まれるドキュメントのみ検索できるようにします。
この仕組みを利用して、Azure OpenAI Service(以下AOAI)とAI Searchの回答生成で参照するドキュメントの制御をすることができます。
下記図の例では、ファイルAにはACLにXXXというoidが設定されています。ユーザーAはoidがXXXであるため、RAGの回答生成時に参照情報としてファイルAを利用できます。
一方でユーザーBはoidがXXXでないため、RAGの回答生成にファイルBを利用することができません。
今回はAOAIによる回答生成は含めず、AI Searchの検索結果がフィルターされているかまでを対象に説明します。
実装
実装の概要としては以下です。
- ADLSにファイルをアップロードし、ACLを設定する
- Pythonでインデックスを作成する
ADLSにファイルをアップロードし、ACLを設定する
ADLSの任意のディレクトリにファイルをアップロードします。
次にACLを設定します。
ADLSはファイルを選択し、右クリックから「ACLの管理」を選択します。
「プリンシパルの追加」をクリックします。
権限を付与したいユーザーを検索し、対象のユーザーにチェックを入れ、「選択を」クリックします。
追加したユーザーに対して、「読み取り」にチェックを入れ、「保存」をクリックします。
以上で、ADLSにアップロードしたファイルにACLを設定することができました。
Pythonでインデックスを作成する
まず、ライブラリをインポートします。
from azure.search.documents import SearchClient
from azure.search.documents.indexes import SearchIndexClient
from azure.search.documents.indexes.models import (
SearchableField,
SearchFieldDataType,
SearchIndex,
SimpleField
)
from azure.core.credentials import AzureKeyCredential
from azure.storage.filedatalake import DataLakeServiceClient
次に変数を設定します。
# AI Searchの設定
searchservice="<AI Searchのサービス名>"
endpoint=f"https://{searchservice}.search.windows.net/"
index_name="<インデックス名>"
credential=AzureKeyCredential("<AI Searchキー>")
# ADLSの設定
data_lake_storage_account="<ADLSのサービス名>"
account_url=f"https://{data_lake_storage_account}.dfs.core.windows.net"
adls_credential="<ADLSのキー>"
次にインデックスを作成します。
# インデックス作成用クライアントの生成
search_index_client = SearchIndexClient(endpoint=endpoint, credential=credential)
# インデックスのフィールド定義とインデックスの名設定
fields = [
SimpleField(name="id", type="Edm.String", key=True),
SearchableField(name="content", type="Edm.String", analyzer_name="ja.microsoft"),
SimpleField(name="oids", type=SearchFieldDataType.Collection(SearchFieldDataType.String), filterable=True),
]
index = SearchIndex(
name=index_name,
fields=fields
)
# インデックスの作成
search_index_client.create_index(index)
ALDSからファイルのACLと中身を読み取ります。
# ADLSのfile clientを取得し、そこからパスで指定したファイルをダウンロードする
service_client = DataLakeServiceClient(account_url=account_url, credential=adls_credential)
file_system_client = service_client.get_file_system_client(file_system="<コンテナー名>")
file_client = file_system_client.get_file_client("<ファイルのパス。この例では/demo/sample.txt>")
download = file_client.download_file()
# ダウンロードしたファイルの中身を取得
file_contents = download.readall()
text = file_contents.decode("utf-8")
# file clientからaclを取得する
access_control = file_client.get_access_control()
acl_list = access_control["acl"]
acl_list = acl_list.split(",")
acls = {"oids": []}
for acl in acl_list:
acl_parts: list = acl.split(":")
if len(acl_parts) != 3:
continue
if len(acl_parts[1]) == 0:
continue
if acl_parts[0] == "user" and "r" in acl_parts[2]:
acls["oids"].append(acl_parts[1])
インデックスにファイルから取得した情報を登録します。
# 登録する情報を定義します。
documents = [
{
"id": "sample_txt",
"content": text,
"oids": acls["oids"],
}
]
# インデックスに登録するためのクライアントを生成します。
search_client=SearchClient(endpoint=endpoint, index_name=index_name, credential=credential)
# インデクッスに登録します。
search_client.upload_documents(documents)
以上で実装は完了です。
検証
検証では、ACLの権限がある想定で検索した場合と、権限がない場合とで検索をします。
権限ある状態での検索
ACLに設定したoidを持つユーザーがシステムを利用した想定で検索をします。
# oidを設定
oid=acls['oids'][0]
# oidsのフィルタリングを用いた検索クエリを実行
results = search_client.search(search_text="*", filter=f"oids/any(o:search.in(o, '{oid}'))")
# 検索数を表示
for result in results:
print(result)
結果は以下の様に登録した情報が出力されます。
権限がない状態での検索
次にoidを登録されていないものにして検索します。
# oidを設定
oid="11111111-1111-1111-1111-111111111111"
# oidsのフィルタリングを用いた検索クエリを実行
results = search_client.search(search_text="*", filter=f"oids/any(o:search.in(o, '{oid}'))")
# 検索数を表示
for result in results:
print(result)
上記では検索にヒットせず、何も表示されません。
まとめ
ADLSのACLとAI Searchを用いた検索フィルターを実装しました。
権限ある資料に対してはAI Searchの検索結果にヒットし、権限のない状態ではヒットしないことが分かりました。
これを利用すれば、RAGの回答生成の参照資料をユーザーの権限により制限することが可能です。