はじめに

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を追加してやれば良さそうです。