メインコンテンツまでスキップ

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

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

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

AI Bridgeを使用したベクトル検索取得ツールをローカルで開発​

Databricks ベクトル検索取得ツールの開発を開始する最も簡単な方法は、ローカルです。databricks-langchain や databricks-openai などの Databricks AI Bridge パッケージを使用して、クエリ パラメーターを使用してエージェントとエクスペリメントに取得機能をすばやく追加します。このアプローチにより、初期開発中の迅速なイテレーションが可能になります。

ローカルツールの準備ができたら、エージェントコードの一部として直接プロダクション化したり、見つけやすさとガバナンスが向上する Unity Catalog 関数に移行したりできますが、一定の制限があります。 Unity Catalog関数を持つベクトル検索 retriever ツールを参照してください。

注記

Databricksの外部でホストされている外部ベクトルインデックスを使用するには、Databricksの外部でホストされているベクトルインデックスを使用したベクトル検索 retrieverを参照してください。

The following code prototypes a retriever tool and binds it to an LLM locally so you can chat with the agent to test its tool-calling behavior.

Install the latest version of databricks-langchain which includes Databricks AI Bridge.

Bash
%pip install --upgrade databricks-langchain

The following example queries a hypothetical vector search index that fetches content from Databricks product documentation.

Provide a clear and descriptive tool_description. The agent LLM uses the tool_description to understand the tool and determine when to invoke the tool.

Python
from databricks_langchain import VectorSearchRetrieverTool, ChatDatabricks

# Initialize the retriever tool.
vs_tool = VectorSearchRetrieverTool(
  index_name="catalog.schema.my_databricks_docs_index",
  tool_name="databricks_docs_retriever",
  tool_description="Retrieves information about Databricks products from official Databricks documentation."
)

# Run a query against the vector search index locally for testing
vs_tool.invoke("Databricks Agent Framework?")

# Bind the retriever tool to your Langchain LLM of choice
llm = ChatDatabricks(endpoint="databricks-meta-llama-3-1-70b-instruct")
llm_with_tools = llm.bind_tools([vs_tool])

# Chat with your LLM to test the tool calling functionality
llm_with_tools.invoke("Based on the Databricks documentation, what is Databricks Agent Framework?")
注記

When initializing the VectorSearchRetrieverTool, the text_column and embedding arguments are required for Delta Sync Indexes with self-managed embeddings and Direct Vector Access Indexes. See options for providing embeddings.

For additional details, see the API docs for VectorSearchRetrieverTool

Python
from databricks_langchain import VectorSearchRetrieverTool
from databricks_langchain import DatabricksEmbeddings

embedding_model = DatabricksEmbeddings(
    endpoint="databricks-bge-large-en",
)

vs_tool = VectorSearchRetrieverTool(
  index_name="catalog.schema.index_name", # Index name in the format 'catalog.schema.index'
  num_results=5, # Max number of documents to return
  columns=["primary_key", "text_column"], # List of columns to include in the search
  filters={"text_column LIKE": "Databricks"}, # Filters to apply to the query
  query_type="ANN", # Query type ("ANN" or "HYBRID").
  tool_name="name of the tool", # Used by the LLM to understand the purpose of the tool
  tool_description="Purpose of the tool", # Used by the LLM to understand the purpose of the tool
  text_column="text_column", # Specify text column for embeddings. Required for direct-access index or delta-sync index with self-managed embeddings.
  embedding=embedding_model # The embedding model. Required for direct-access index or delta-sync index with self-managed embeddings.
)

Unity Catalog関数を用いたベクトル検索レトリーバーツール​

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

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

注記

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

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

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
  )

この取得ツールを AI エージェントで使用するには、 UCFunctionToolkitでラップします。 これにより、MLflow による自動トレースが可能になります。

MLflow トレースは、生成 AI アプリケーションの詳細な実行情報をキャプチャします。 各ステップの入力、出力、メタデータをログに記録し、問題のデバッグとパフォーマンスの分析を支援します。

UCFunctionToolkitを使用すると、レトリーバーは、その出力が MLflow レトリーバー スキーマに準拠している場合、MLflow ログに RETRIEVER スパン型を自動的に生成します。「MLflow トレース スキーマ」を参照してください。

UCFunctionToolkitの詳細については Unity Catalog ドキュメントを参照してください。

Python
from unitycatalog.ai.langchain.toolkit import UCFunctionToolkit

toolkit = UCFunctionToolkit(
    function_names=[
        "main.default.databricks_docs_vector_search"
    ]
)
tools = toolkit.tools

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

  • SQL クライアントでは、返される行数またはバイト数の制限がある場合があります。 データの切り捨てを防ぐには、UDF によって返される列値を切り捨てる必要があります。 たとえば、 substring(chunked_text, 0, 8192) を使用して、大きなコンテンツ列のサイズを縮小し、実行中の行の切り捨てを回避できます。
  • このツールは vector_search() 関数のラッパーであるため、 vector_search() 関数と同じ制限が適用されます。 制限事項を参照してください。

Databricks の外部でホストされているベクトルインデックスを使用したベクトル検索 レトリーバー​

ベクトルインデックスが Databricksの外部でホストされている場合は、 Unity Catalogの接続 を作成して外部サービスに接続し、接続エージェントコードを使用できます。 詳細については、「AI エージェント ツールを外部サービスに接続する」を参照してください。

次の例では、Databricks の外部でホストされている PyFunc フレーバー エージェントのベクトル インデックスを呼び出すベクトル検索レトリーバーを作成します。

  1. 外部サービス (この場合は Azure) への Unity Catalog 接続を作成します。

    SQL
    CREATE CONNECTION ${connection_name}
    TYPE HTTP
    OPTIONS (
      host 'https://example.search.windows.net',
      base_path '/',
      bearer_token secret ('<secret-scope>','<secret-key>')
    );

  2. エージェント コードで取得ツールを定義するには、作成した Unity Catalog の接続 を使用します。 この例では、MLflow デコレーターを使用してエージェント トレースを有効にします。

注記

MLflow 取得コンポーネントに準拠するために、取得コンポーネント関数は Document 型を返し、Document クラスの metadata フィールドを使用して、返されたドキュメントに like doc_uri や similarity_score.

Python
import mlflow
import json

from mlflow.entities import Document
from typing import List, Dict, Any
from dataclasses import asdict

class VectorSearchRetriever:
  """
  Class using Databricks Vector Search to retrieve relevant documents.
  """

  def __init__(self):
    self.azure_search_index = "hotels_vector_index"

  @mlflow.trace(span_type="RETRIEVER", name="vector_search")
  def __call__(self, query_vector: List[Any], score_threshold=None) -> List[Document]:
    """
    Performs vector search to retrieve relevant chunks.
    Args:
      query: Search query.
      score_threshold: Score threshold to use for the query.

    Returns:
      List of retrieved Documents.
    """
    from databricks.sdk import WorkspaceClient
    from databricks.sdk.service.serving import ExternalFunctionRequestHttpMethod

    json = {
      "count": true,
      "select": "HotelId, HotelName, Description, Category",
      "vectorQueries": [
        {
          "vector": query_vector,
          "k": 7,
          "fields": "DescriptionVector",
          "kind": "vector",
          "exhaustive": true,
        }
      ],
    }

    response = (
      WorkspaceClient()
      .serving_endpoints.http_request(
        conn=connection_name,
        method=ExternalFunctionRequestHttpMethod.POST,
        path=f"indexes/{self.azure_search_index}/docs/search?api-version=2023-07-01-Preview",
        json=json,
      )
      .text
    )

    documents = self.convert_vector_search_to_documents(response, 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 = []

    for item in vs_results.get("value", []):
      score = item.get("@search.score", 0)

      if score >= score_threshold:
        metadata = {
          "score": score,
          "HotelName": item.get("HotelName"),
          "Category": item.get("Category"),
        }

        doc = Document(
          page_content=item.get("Description", ""),
          metadata=metadata,
          id=item.get("HotelId"),
        )
        docs.append(doc)

    return docs

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

    Python
    retriever = VectorSearchRetriever()
    query = [0.01944167, 0.0040178085 . . .  TRIMMED FOR BREVITY 010858015, -0.017496133]
    results = retriever(query, score_threshold=0.1)

レトリーバースキーマの設定​

レトリーバーまたは span_type="RETRIEVER" から返されたトレースが MLflow の標準レトリーバー スキーマに準拠していない場合は、返されたスキーマを MLflow の予期されるフィールドに手動でマップする必要があります。 これにより、MLflow はレトリーバーを適切にトレースし、ダウンストリーム アプリケーションでトレースを正しくレンダリングできます。

レトリーバースキーマを手動で設定するには、 mlflow.models.set_retriever_schema エージェントを定義するとき。 set_retriever_schema を使用して、返されたテーブルの列名を MLflow の想定フィールド (primary_key、text_column、doc_uriなど) にマップします。

Python
# 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 の自動ロギングを使用できます。 自動トレースを参照してください。

Python
...
@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名が一致している必要があります。

次のステップ​

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

に最終更新