こんにちは！ACSD松崎です。

Llama_index関連のブログ投稿の第三弾として、今回は、SimpleDirectoryReaderとGPTVectorStoreIndexを使ってファイルベースのFAQエンジンを作る際に回答に使われたドキュメントの情報を生成＆回答に付与する方法を解説します。

概要

Llama_indexの SimpleDirectoryReader を使う際のサンプルコードは、

documents = SimpleDirectoryReader("data").load_data()

のような例が多いと思います。代わりにSimpleDirectoryReader 生成時の第2引数にファイル名をメタデータとして取得するラムダ式を渡すことで、後段でqueryをする際のRESPONSEオブジェクトに回答に使われたファイル名が含まれるようになります。

file_metadata = lambda x: {"filename": x}
documents = SimpleDirectoryReader("data", file_metadata=file_metadata).load_data()

参考：

以下は、実際にqueryを行う際に、ファイル名を取得している例になります。ファイル名のメタデータは response.source_nodes 内の node.node.metadata['filename'] に含まれるため、それをprintして出力しています。

question = "コロンビアについて教えて"
response = query_engine.query(question)

for node in response.source_nodes:
        print("------------------------------------------------------")
        print("metadata ", node.node.metadata['filename'])
        print("------------------------------------------------------")

上記のコードの出力例は以下となります。

------------------------------------------------------
metadata  data/コロンビア.data
------------------------------------------------------
------------------------------------------------------
metadata  data/インフルエンザ.data
------------------------------------------------------

※ Llama_indexの以前のverではメタデータとしてのファイル名は node.node.extrainfo に含まれていたようですが、v0.6.34のタイミングで node.node.metadata に引っ越しとなったようです。

参考：

コード全文

以下はGoogle Colab上で、実行をした例になります。動作させるに際し、Google Colab上に data ディレクトリを作成し、textファイルを格納する必要があるのでご注意ください。

!pip install openai
!pip install llama-index

!pip freeze | grep openai
!pip freeze | grep llama-index

import openai
openai.api_key = "sk-xxx"

from llama_index import GPTVectorStoreIndex, SimpleDirectoryReader

# documents作成
file_metadata = lambda x: {"filename": x}
documents = SimpleDirectoryReader("data", file_metadata=file_metadata).load_data()

# documentsの中身を確認
for doc in documents:
    print(doc)

# index作成
index = GPTVectorStoreIndex.from_documents(documents)

# query_engine作成
query_engine = index.as_query_engine()

# 質問の実行
question = "コロンビアについて教えて"
response = query_engine.query(question)

# responseの解析
for node in response.source_nodes:
        print("------------------------------------------------------")
        print("score ", node.score)
        print("id_", node.id_)
        print("metadata ", node.node.metadata['filename'])
        print("------------------------------------------------------")