ベクトル検索レトリーバーを作成する
プレビュー
この機能はパブリックプレビュー段階です。
Mosaic AI Agent Framework を使用してレトリーバーを作成する方法を学びます。 レトリーバーは、ベクトル検索インデックスを使用して関連ドキュメントを検索して返す エージェントツール の一種です。 レトリーバーは、RAG(Retrieval Augmented Generation)アプリケーションのコアコンポーネントです。
要件
MLflow
Document
は、MLflow バージョン 2.14.0 以降でのみ使用できます。既存の ベクトル検索インデックス。
PyFunc レトリーバの例
次の例では、 databricks-vectorsearch を使用して、 フィルターを使用したベクトル検索の類似性検索を実行する基本的な取得ツールを作成します。 MLflow デコレーターを使用して 、エージェント トレースを有効にします。
レトリーバー 関数は 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
レトリーバー を実行するには、次の 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_schema名が一致している必要があります。
ベクトル検索結果のフィルタリング
ベクトル検索フィルタを使用して、検索範囲をデータのサブセットに制限できます。
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 リポジトリ を参照してください。
LangChain レトリーバーの例: デモ: Mosaic AI エージェント フレームワークとエージェントの評価
PyFunc レトリーバーの例: Agent アプリのサンプルコード