非構造化検索AIエージェントツール

プレビュー

この記事では、Mosaic AI Agent Framework を使用して非構造化データを取得するための AI エージェントツールを作成する方法について説明します。非構造化リトリーバーを使用すると、エージェントはベクトル検索インデックスを使用して、ドキュメントコーパスなどの非構造化データソースをクエリできます。

エージェントツールの詳細については、「 AI エージェントツールを作成する」を参照してください。

Unity Catalog機能を持つベクトル検索レトリーバーツール

次の例では、 Mosaic AI Vector Search インデックスからデータをクエリできる取得ツール用の Unity Catalog 関数を作成します。

Unity Catalog 関数は、Databricksドキュメントを含む架空のベクトル検索インデックスをクエリdatabricks_docs_vector_search。Databricks SQL 関数 vector_search() をラップし、エイリアス page_content と metadata を使用して、その出力を MLflow 取得者スキーマと一致させます。

注：

MLflow 取得者スキーマに準拠するには、最上位の出力キーとしてではなく、SQL マップ関数を使用して、追加のメタデータ列を metadata 列に追加する必要があります。

ノートブックまたは SQL エディターで次のコードを実行します。

CREATE OR REPLACE FUNCTION main.default.databricks_docs_vector_search (
  -- The agent uses this comment to determine how to generate the query string parameter.
  query STRING
  COMMENT 'The query string for searching Databricks documentation.'
) RETURNS TABLE
-- The agent uses this comment to determine when to call this tool. It describes the types of documents and information contained within the index.
COMMENT 'Executes a search on Databricks documentation to retrieve text documents most relevant to the input query.' RETURN
SELECT
  chunked_text as page_content,
  map('doc_uri', url, 'chunk_id', chunk_id) as metadata
FROM
  vector_search(
    -- Specify your Vector Search index name here
    index => 'catalog.schema.databricks_docs_index',
    query => query,
    num_results => 5
  )

このレトリーバーツールには、次の注意事項があります。

MLflow は、この Unity Catalog 関数を RETRIEVER スパンタイプではなく TOOL スパンタイプとしてトレースします。その結果、エージェントレビューアプリやAI PlaygroundなどのダウンストリームのAgent Frameworkアプリケーションには、チャンクへのリンクなど、レトリーバー固有の詳細が表示されません。スパンタイプの詳細については、「 MLflow トレーシングスキーマ」を参照してください。
SQL クライアントでは、返される行数またはバイト数の制限がある場合があります。データの切り捨てを防ぐには、UDF によって返される列値を切り捨てる必要があります。たとえば、 substring(chunked_text, 0, 8192) を使用して、大きなコンテンツ列のサイズを縮小し、実行中の行の切り捨てを回避できます。
このツールは vector_search() 関数のラッパーであるため、 vector_search() 関数と同じ制限が適用されます。制限事項を参照してください。

この例がユースケースに適していない場合は、代わりにカスタムエージェントコードを使用してベクトル検索取得ツールを作成してください。

エージェントコード付きベクトル検索レトリーバ(PyFunc)

次の例では、エージェントコードで PyFunc フレーバーのエージェントのベクトル検索レトリーバーを作成します。

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

注：

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"],
)

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

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

エージェント作成時に設定された取得者スキーマは、レビューアプリや評価セットなどのダウンストリームアプリケーションやワークフローに影響します。具体的には、 doc_uri 列は、取得者によって返されるドキュメントのプライマリ識別子として機能します。

レビューアプリにはdoc_uriが表示され、レビュー担当者が回答を評価し、ドキュメントの出所を追跡するのに役立ちます。「アプリのUIを確認する」を参照してください。
評価セット では、 doc_uri を使用して、レトリーバーの結果を事前定義された評価データセットと比較し、レトリーバーの再現率と精度を判断します。評価セットを参照してください。

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

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 条件を満たすドキュメントが含まれます。

次のステップ

Unity Catalog 関数エージェントツールを作成したら、そのツールを AI エージェントに追加します。「エージェントへの Unity Catalog ツールの追加」を参照してください。