ベクトル検索取得ツールの作成

プレビュー

この機能はパブリックプレビュー段階です。

Mosaic AI Agent Framework を使用してレトリーバーを作成する方法を学びます。 レトリーバーは、ベクトル検索インデックスを使用して関連ドキュメントを検索して返す エージェントツール の一種です。 レトリーバーは、RAG(Retrieval Augmented Generation)アプリケーションのコアコンポーネントです。

要件

PyFunc レトリーバの例

次の例では、 databricks-vectorsearch を使用して、 フィルターを使用したベクトル検索の類似性検索を実行する基本的な取得ツールを作成します。 MLflow デコレーターを使用して 、エージェント トレースを有効にします。

retriever 関数は Document 型を返す必要があります。 Document クラスの metadata フィールドを使用して、返されたドキュメントに like doc_uri や similarity_scoreなどの属性を追加します。

エージェント モジュールまたはエージェント ノートブックで次のコードを使用します。

import mlflow
import json

from mlflow.entities import Document
from typing import List, Dict, Any
from dataclasses import asdict
from databricks.vector_search.client import VectorSearchClient

class VectorSearchRetriever:
    """
    Class using Databricks Vector Search to retrieve relevant documents.
    """
    def __init__(self):
        self.vector_search_client = VectorSearchClient(disable_notice=True)
        # TODO: Replace this with the list of column names to return in the result when querying Vector Search
        self.columns = ["chunk_id", "text_column", "doc_uri"]
        self.vector_search_index = self.vector_search_client.get_index(
            index_name="catalog.schema.chunked_docs_index"
        )
        mlflow.models.set_retriever_schema(
            name="vector_search",
            primary_key="chunk_id",
            text_column="text_column",
            doc_uri="doc_uri"
        )

    @mlflow.trace(span_type="RETRIEVER", name="vector_search")
    def __call__(
        self,
        query: str,
        filters: Dict[Any, Any] = None,
        score_threshold = None
    ) -> List[Document]:
        """
        Performs vector search to retrieve relevant chunks.
        Args:
            query: Search query.
            filters: Optional filters to apply to the search. Filters must follow the Databricks Vector Search filter spec
            score_threshold: Score threshold to use for the query.

        Returns:
            List of retrieved Documents.
        """

        results = self.vector_search_index.similarity_search(
            query_text=query,
            columns=self.columns,
            filters=filters,
            num_results=5,
            query_type="ann"
        )

        documents = self.convert_vector_search_to_documents(
            results, score_threshold
        )
        return [asdict(doc) for doc in documents]

    @mlflow.trace(span_type="PARSER")
    def convert_vector_search_to_documents(
        self, vs_results, score_threshold
    ) -> List[Document]:

        docs = []
        column_names = [column["name"] for column in vs_results.get("manifest", {}).get("columns", [])]
        result_row_count = vs_results.get("result", {}).get("row_count", 0)

        if result_row_count > 0:
            for item in vs_results["result"]["data_array"]:
                metadata = {}
                score = item[-1]

                if score >= score_threshold:
                    metadata["similarity_score"] = score
                    for i, field in enumerate(item[:-1]):
                        metadata[column_names[i]] = field

                    page_content = metadata.pop("text_column", None)

                    if page_content:
                        doc = Document(
                            page_content=page_content,
                            metadata=metadata
                        )
                        docs.append(doc)

        return docs

retriever を実行するには、次の Python コードを実行します。 必要に応じて、結果をフィルタリングするために、要求に ベクトル検索フィルター を含めることができます。

retriever = VectorSearchRetriever()
query = "What is Databricks?"
filters={"text_column LIKE": "Databricks"},
results = retriever(query, filters=filters, score_threshold=0.1)

取得者スキーマの設定

取得子が正しくトレースされていることを確認するには、 mlflow.models.set_retriever_schema エージェントをコードで定義する場合。 set_retriever_schema を使用して、返されたテーブルの列名を MLflow の想定フィールド (primary_key、text_column、doc_uriなど) にマップします。

# Define the retriever's schema by providing your column names
mlflow.models.set_retriever_schema(
    name="vector_search",
    primary_key="chunk_id",
    text_column="text_column",
    doc_uri="doc_uri"
    # other_columns=["column1", "column2"],
)

注：

doc_uri列は、レトリーバーのパフォーマンスを評価するときに特に重要です。doc_uri は、レトリーバーによって返されるドキュメントの主な識別子であり、それらをグラウンド トゥルース評価セットと比較できます。 評価セットを参照してください。

また、レトリーバーのスキーマで追加の列を指定するには、 other_columns フィールドに列名のリストを指定します。

複数のレトリーバーがある場合は、各レトリーバースキーマに一意の名前を使用して、複数のスキーマを定義できます。

レトリーバーをトレースする

MLflow トレースは、エージェントの実行に関する詳細な情報をキャプチャすることで、可観測性を追加します。 これにより、要求の各中間ステップに関連付けられた入力、出力、メタデータを記録する方法が提供され、バグや予期しない動作のソースを簡単に特定できます。

この例では、 @mlflow.trace デコレータ を使用して、レトリーバとパーサのトレースを作成します。 トレース方法を設定するためのその他のオプションについては、「 エージェントの MLflow トレース」を参照してください。

デコレータは、関数が呼び出されたときに開始し、関数が戻ったときに終了する スパン を作成します。 MLflow は、関数の入力と出力、および関数から発生した例外を自動的に記録します。

注：

LangChain、LlamaIndex、および OpenAI ライブラリのユーザーは、デコレータでトレースを手動で定義する代わりに、MLflow 自動ロギングを使用できます。 「自動ログを使用してエージェントにトレースを追加する」を参照してください。

...
@mlflow.trace(span_type="RETRIEVER", name="vector_search")
def __call__(self, query: str) -> List[Document]:
  ...

Agent Evaluation や AI Playground などのダウンストリーム アプリケーションでレトリーバー トレースが正しくレンダリングされるようにするには、デコレータが次の要件を満たしていることを確認してください。

  • span_type="RETRIEVER"を使用して、関数がオブジェクトを返すList[Document]を確認します。「Retriever スパン」を参照してください。

  • トレースを正しく構成するには、トレース名とretriever_schema名が一致している必要があります。

Filter ベクトル検索の結果

ベクトル検索フィルタを使用して、検索範囲をデータのサブセットに制限できます。

VectorSearchRetriever の filters パラメーターは、 Databricks ベクトル検索フィルター仕様を使用してフィルター条件を定義します。

filters = {"text_column LIKE": "Databricks"}

__call__メソッド内では、filtersディクショナリが直接similarity_search関数に渡されます。

results = self.vector_search_index.similarity_search(
    query_text=query,
    columns=self.columns,
    filters=filters,
    num_results=5,
    query_type="ann"
)

初期フィルター処理の後、 score_threshold パラメーターは、最小類似性スコアを設定することで追加のフィルター処理を提供します。

if score >= score_threshold:
    metadata["similarity_score"] = score

最終結果には、 filters 条件と score_threshold 条件を満たすドキュメントが含まれます。

レトリーバー・サンプル・アプリケーション

レトリーバーを使用する AI エージェントの例については、 genai-cookbook GitHub リポジトリ を参照してください。