はじめに
langchainのAgentExecutorの出力に、（中間ステップで実行した）ベクトル検索のスコアを含める方法を検討しました。
動機
langchainを使った検索UIを作るうえで、ベクトル検索の結果を事後に補正したくなる時があります。langchainのAgent機能（AgentExecutor）を使うと、ユーザ入力を言語モデルで検索ワードに変換し、検索結果をえるという処理を簡潔に書けて便利です。しかし、検索結果の文章は中間出力として取得できるのですが、検索スコアは含まれていませんでした。そこで、AgentExecutorの出力に検索スコアを含める方法を検討しました。
環境
% pip list | grep langchain
langchain 0.0.277
実装
結論としては、VectorStoreクラスのsimilarity_search_with_score系の関数をオーバーライドして、Documentのmetadataにscoreを書き込むようにするのが簡単だと思いました。
例えば、FAISSクラスの場合、以下のようなクラスを定義してToolに渡すことで、実現できました。
from langchain.vectorstores.faiss import FAISS
import operator
import os
import pickle
import uuid
import warnings
from pathlib import Path
from typing import (
Any,
Callable,
Dict,
Iterable,
List,
Optional,
Sized,
Tuple,
)
import numpy as np
from langchain.docstore.document import Document
from langchain.vectorstores.utils import DistanceStrategy
from langchain.vectorstores.faiss import dependable_faiss_import
class MyFAISS(FAISS):
def similarity_search_with_score_by_vector(
self,
embedding: List[float],
k: int = 4,
filter: Optional[Dict[str, Any]] = None,
fetch_k: int = 20,
**kwargs: Any,
) -> List[Tuple[Document, float]]:
"""Return docs most similar to query.
Args:
embedding: Embedding vector to look up documents similar to.
k: Number of Documents to return. Defaults to 4.
filter (Optional[Dict[str, Any]]): Filter by metadata. Defaults to None.
fetch_k: (Optional[int]) Number of Documents to fetch before filtering.
Defaults to 20.
**kwargs: kwargs to be passed to similarity search. Can include:
score_threshold: Optional, a floating point value between 0 to 1 to
filter the resulting set of retrieved docs
Returns:
List of documents most similar to the query text and L2 distance
in float for each. Lower score represents more similarity.
"""
faiss = dependable_faiss_import()
vector = np.array([embedding], dtype=np.float32)
if self._normalize_L2:
faiss.normalize_L2(vector)
scores, indices = self.index.search(vector, k if filter is None else fetch_k)
docs = []
for j, i in enumerate(indices[0]):
if i == -1:
# This happens when not enough docs are returned.
continue
_id = self.index_to_docstore_id[i]
doc = self.docstore.search(_id)
# オリジナル版に追記。metadataにscoreを保存
doc.metadata["score"] = scores[0][j]
# 追記ここまで
if not isinstance(doc, Document):
raise ValueError(f"Could not find document for id {_id}, got {doc}")
if filter is not None:
filter = {
key: [value] if not isinstance(value, list) else value
for key, value in filter.items()
}
if all(doc.metadata.get(key) in value for key, value in filter.items()):
docs.append((doc, scores[0][j]))
else:
docs.append((doc, scores[0][j]))
score_threshold = kwargs.get("score_threshold")
if score_threshold is not None:
cmp = (
operator.ge
if self.distance_strategy
in (DistanceStrategy.MAX_INNER_PRODUCT, DistanceStrategy.JACCARD)
else operator.le
)
docs = [
(doc, similarity)
for doc, similarity in docs
if cmp(similarity, score_threshold)
]
return docs[:k]
# 既存のFAISSクラスを変換するためのクラスメソッドを追加
@classmethod
def from_vectorstore(cls, vectorstore: FAISS):
"""
オリジナルから必要な変数をコピーしてクラスを初期化。以下参考。
Initialize with necessary components.
self.embedding_function = embedding_function
self.index = index
self.docstore = docstore
self.index_to_docstore_id = index_to_docstore_id
self.distance_strategy = distance_strategy
self.override_relevance_score_fn = relevance_score_fn
self._normalize_L2 = normalize_L2
"""
return cls(
vectorstore.embedding_function
, vectorstore.index
, vectorstore.docstore
, vectorstore.index_to_docstore_id
, vectorstore.override_relevance_score_fn
, vectorstore._normalize_L2
)
以下のように使います。
from langchain.chat_models import ChatOpenAI
from langchain.agents import initialize_agent, AgentType, Tool
from langchain.embeddings import OpenAIEmbeddings
vectorstore = MyFAISS.from_texts(["「こんにちは」は昼の挨拶です。", "「さようなら」は別れの挨拶です。"], OpenAIEmbeddings())
retriever = vectorstore.as_retriever(
search_kwargs=dict(k=3)
)
tools = [
Tool(
name = "greeting_search"
, func = retriever.get_relevant_documents
, description = "挨拶を検索したいときに使う"
)
]
llm = ChatOpenAI(model="gpt-3.5-turbo")
agent = initialize_agent(tools, llm, AgentType.OPENAI_FUNCTIONS, return_intermediate_steps=True)
res = agent.invoke({"input":"別れの挨拶は？"})
# 出力の取得
output = res["output"]
# 検索結果の取得
intermediate_steps = res["intermediate_steps"] # 中間出力
print(output)
print(intermediate_steps)
別れの挨拶としては「さようなら」が一般的です。
[(_FunctionsAgentAction(tool='greeting_search', tool_input='別れ', log='\nInvoking: `greeting_search` with `別れ`\n\n\n', message_log=[AIMessage(content='', additional_kwargs={'function_call': {'name': 'greeting_search', 'arguments': '{\n "__arg1": "別れ"\n}'}}, example=False)]), [Document(page_content='「さようなら」は別れの挨拶です。', metadata={'score': 0.26613048}), Document(page_content='「こんにちは」は昼の挨拶です。', metadata={'score': 0.36656606})])]
intermediate_stepsに含まれる検索結果（Document）に、中に、scoreを含むDocumentのデータが存在していました。
解説
AgentExecutorの基本的な仕組み
initialize_agentの戻り値はAgentExecutorです。AgentExecutorはToolを使ってユーザの入力に回答する一連の処理を行うためのクラスです。AgentExecutorの初期化時にreturn_intermediate_steps=Trueを指定することで、中間ステップの出力も取得できます。中間ステップの出力とはToolの実行結果です。Toolの実行結果とは、Toolの初期化時に渡した関数の戻り値です。
tools = [
Tool(
name = "greeting_search"
, func = retriever.get_relevant_documents
, description = "挨拶を検索したいときに使う"
)
]
上記で言うと、retriever.get_relevant_documentsの戻り値がintermediate_stepsに含まれます。
検討したこと
要するに、Toolに与える関数がスコア込みの検索結果を返すものであれば、intermediate_steps経由でAgentExecutorの出力としてscoreを取り出せることになります。
get_relevant_documentsを全く別の関数にしてしまうことも検討したのですが、get_relevant_documents含めBaseRetrieverはToolとして使いやすい工夫がいろいろなされたクラスなので、なるべく活かしたいと思いました。AgentExecutor側が問題なく動くことを保証するためにも、Toolの戻り値はDocumentのリストのままにしておいた方が良いと思いました。そこで、Documentのmetadataにscoreを追加することにしました。
ソースコードを読むと、VectorStoreの中でもFAISSの場合は、similarity_search_with_score_by_vectorという関数で検索結果をDocument化しています。同じ関数の中でscoreも得ていますので、この関数だけoverrideすれば、目的が達成できるとわかりました。
FAISS以外のVectorStoreの場合、similarity_search_with_scoreが存在しているとは限りませんが、出力をDocument化している関数が何かしらあるはずなので、そこでscoreを追加してやれば良さそうです。