[LangChain]FAISSベクトルストア検索でメタデータを活用する

はじめに

LangChainのFAISSベクトルストア検索でメタデータを使った検索方法についてまとめました。

実装例

"おはようございます"、"こんにちは"、"こんばんは"という日本語と英語のテキストをサンプルデータとして使用します。
metadataには日本語か英語かという情報が付与されています。

メタデータを使った類似度検索

from pprint import pprint
from langchain.vectorstores import FAISS
from langchain.embeddings import OpenAIEmbeddings
from langchain.docstore.document import Document

chunk_size = 500
embeddings = OpenAIEmbeddings(chunk_size=chunk_size)

documents = [
    Document(page_content="おはようございます", metadata=dict(lang="ja")),
    Document(page_content="こんにちは", metadata=dict(lang="ja")),
    Document(page_content="こんばんは", metadata=dict(lang="ja")),
    Document(page_content="Good morning", metadata=dict(lang="en")),
    Document(page_content="Good afternoon", metadata=dict(lang="en")),
    Document(page_content="Good evening", metadata=dict(lang="en")),
]

faiss = FAISS.from_documents(documents, embeddings)
docs = faiss.similarity_search("こんにちは", k=1)
pprint(docs)

"こんにちは"で検索すると日本語の"こんにちは"がヒットします。

[Document(page_content='こんにちは', metadata={'lang': "ja"})]

検索対象を英語に限定してみます。
similarity_searchのfilterに検索対象をdict型で与えます。今回は英語のみを対象に検索したいので、lang="en"としています。

docs = faiss.similarity_search("こんにちは", k=1, filter=dict(lang="en"))

"こんにちは"ではなく"Good afternoon"がヒットしました。

[Document(page_content='Good afternoon', metadata={'lang': 'en'})]

dict(hoge="", fuga="")のようにするとAND検索ができますが、OR検索はできないようです。

対象の文章を全て取得する

最後に対象のメタデータが付与された文章を全て取得する方法を紹介します。
エンべディングされたデータはdocstore._dictに格納されています。
このデータに対してフィルタリング処理を行えば、必要な文章を全て取得することができます。

def find_by_metadata(
    faiss: FAISS, filter: dict[str, Any] | None = None
) -> dict[str, Document]:
    if filter is None:
        return faiss.docstore._dict

    response = {
        key: item
        for key, item in faiss.docstore._dict.items()
        if all(item.metadata.get(k) == value for k, value in filter.items())
    }

    return response


docs = find_by_metadata(faiss, dict(lang="en"))
pprint(docs)

{'54883d65-fef7-44d3-b00b-fd56d8bd52d0': Document(page_content='Good morning', metadata={'lang': 'en'}),
 'f7ea9f81-044b-4f3d-b17b-f46a7e3c3058': Document(page_content='Good afternoon', metadata={'lang': 'en'}),
 'fbdd669d-b295-4a43-94d3-182823201b58': Document(page_content='Good evening', metadata={'lang': 'en'})}

このようにメタデータが英語の文章を全て取得することができました！