メインコンテンツまでスキップ
に最終更新

エージェントを非構造化データに接続する

AI エージェントは、質問に答えたり、コンテキストに応じた応答を提供したりするために、ドキュメント コレクション、ナレッジ ベース、テキスト コーパスなどの非構造化データを照会する必要があることがよくあります。

Databricks 、エージェントを地下鉄検索インデックスや外部ベクトル ストアの非構造化データに接続するための複数のアプローチを提供します。 事前構成された MCP サーバーを使用してDatabricks一斉検索インデックスに即座にアクセスしたり、 AI Bridge パッケージを使用してローカルで取得ツールを開発したり、特殊なワークフロー用のカスタム取得関数を構築したりできます。

MCP を使用してDatabricks自動検索インデックスをクエリする​

エージェントがDatabricks地下鉄検索インデックスをクエリする必要がある場合は、 Databricksで管理される MCP サーバーを使用します。 開始する前に、Databricksが管理する埋め込みデータを使用してベクトル検索インデックスを作成してください。ベクトル検索のエンドポイントとインデックスの作成を参照してください。

待機検索用の管理対象 MCP URL は: https://<workspace-hostname>/api/2.0/mcp/vector-search/{catalog}/{schema}/{index_name}です。

以下の例は、エージェントをベクトル検索インデックスに接続する方法を示しています。<catalog> 、 <schema> 、 <index-name>ご使用のベクトル検索インデックスの名前に置き換えてください。

Python
from agents import Agent, Runner
from databricks.sdk import WorkspaceClient
from databricks_openai.agents import McpServer

workspace_client = WorkspaceClient()

async with McpServer.from_vector_search(
    catalog="<catalog>",
    schema="<schema>",
    index_name="<index-name>",
    workspace_client=workspace_client,
    name="vector-search",
) as vs_server:
    agent = Agent(
        name="Research assistant",
        instructions="You are a research assistant. Use the vector search tool to find relevant documents and answer questions.",
        model="databricks-claude-sonnet-4-5",
        mcp_servers=[vs_server],
    )
    result = await Runner.run(agent, "What is the return policy?")
    print(result.final_output)

アプリにdatabricks.ymlのベクトル検索インデックスへのアクセス権限を付与します。

YAML
resources:
  apps:
    my_agent_app:
      resources:
        - name: 'my_vector_index'
          uc_securable:
            securable_full_name: '<catalog>.<schema>.<index-name>'
            securable_type: 'TABLE'
            permission: 'SELECT'

その他のアプローチ​

Databricks 外のベクトル検索インデックスをクエリする

Databricks以外でホストされているベクトル検索インデックスをクエリする​

ベクターインデックスがDatabricks以外でホストされている場合は、 Unity Catalog接続を作成して外部サービスに接続し、その接続をエージェントコードで使用できます。 エージェントを外部サービスに接続する方法については、こちらをご覧ください。

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

  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 取得コンポーネントに準拠するには、取得コンポーネントで List[Document] オブジェクトを返し、Document クラスの metadata フィールドを使用して、返されたドキュメントに doc_uri や similarity_scoreなどの属性を追加する必要があります。「MLflow ドキュメント」を参照してください。

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.
    """
    import requests
    from databricks.sdk import WorkspaceClient

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

    response = requests.post(
      f"{w.config.host}/api/2.0/unity-catalog/connections/{connection_name}/proxy/indexes/{self.azure_search_index}/docs/search?api-version=2023-07-01-Preview",
      headers={
        **w.config.authenticate(),
        "Content-Type": "application/json",
      },
      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)

地元のレトリバーを育成する

AI Bridgeを使用してローカルでレトリバーを開発する​

Databricks高速検索取得ツールをローカルで構築するには、 databricks-langchainやdatabricks-openaiなどのDatabricks AI Bridge パッケージを使用します。 これらのパッケージには、既存の Databricks リソースからリトリーバーを作成するためのfrom_vector_searchやfrom_uc_functionなどのヘルパー関数が含まれています。

Databricks AI Bridgeを含む最新バージョンのdatabricks-langchainをインストールしてください。

Bash
%pip install --upgrade databricks-langchain

次のコードは、架空のベクトル検索インデックスをクエリし、それをローカルの LLM にバインドして、ツール呼び出しの動作をテストできるようにする取得ツールのプロトタイプを作成します。

エージェントがツールを理解し、いつ呼び出すかを判断するのに役立つ説明的な tool_description を提供します。

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-claude-sonnet-4-5")
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?")

直接アクセス インデックスを使用するか、自己管理型の埋め込みを使用する Delta Sync インデックスを使用するシナリオでは、 VectorSearchRetrieverTool を構成し、カスタム埋め込みモデルとテキスト列を指定する必要があります。埋め込みを提供するためのオプションを参照してください。

次の例は、columns キーと embedding キーを使用して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.
)

詳細については、 VectorSearchRetrieverToolの API ドキュメントを参照してください。

ローカルツールが準備できたら、エージェントコードの一部として直接本番環境に導入することも、 Unity Catalog機能に移行することもできます。Unityカタログ機能に移行すると、発見性やガバナンスが向上しますが、いくつかの制限があります。

UC 関数を使用してDatabricksクエリする (非推奨)

UC 関数を使用してDatabricksクエリする (非推奨)​

注記

Databricks recommends MCP servers for most agent tools, but defining tools with Unity Catalog functions remains available for prototyping.

Mosaic AI Vector Search インデックス クエリをラップする Unity Catalog 関数を作成できます。このアプローチでは、次のことを行います。

  • ガバナンスと発見可能性を備えた本番運用のユースケースをサポート
  • 内部で vector_search() SQL 関数を使用します
  • 自動 MLflow トレースをサポート
    • 関数の出力を MLflow 取得コンポーネント に揃えるには、 page_content エイリアスと metadata エイリアスを使用する必要があります。
    • 追加のメタデータ列は、最上位の出力キーとしてではなく、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 ログに RETRIEVER スパンの種類を自動的に生成することで、MLflow による自動トレースが可能になります。

Python
from unitycatalog.ai.langchain.toolkit import UCFunctionToolkit

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

Unity Catalog 取得ツールには、次の注意事項があります。

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

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

取得ツールにトレース機能を追加する​

MLflow トレースを追加して、レトリーバーを監視およびデバッグします。トレースでは、実行の各ステップの入力、出力、およびメタデータを表示できます。

前の例では、__call__メソッドと解析メソッドの両方に @mlflow.trace デコレータを追加しています。デコレータは、関数が呼び出されたときに開始し、関数が戻ったときに終了する スパン を作成します。MLflow は、関数の入力と出力、および発生した例外を自動的に記録します。

注記

LangChain、LlamaIndex、OpenAI ライブラリのユーザーは、デコレータを使用してトレースを手動で定義するだけでなく、MLflow 自動ログ記録も使用できます。「アプリケーションへのトレースの追加: 自動トレースと手動トレース」を参照してください。

Python
import mlflow
from mlflow.entities import Document

# This code snippet has been truncated for brevity. See the full retriever example above.
class VectorSearchRetriever:
  ...

  # Create a RETRIEVER span. The span name must match the retriever schema name.
  @mlflow.trace(span_type="RETRIEVER", name="vector_search")
  def __call__(...) -> List[Document]:
    ...

  # Create a PARSER span.
  @mlflow.trace(span_type="PARSER")
  def parse_results(...) -> List[Document]:
    ...

エージェント評価やAI Playgroundなどの下流アプリケーションがリトリーバートレースを正しくレンダリングすることを確認するには、デコレータが以下の要件を満たしていることを確認してください。

  • MLflowのリトリーバーのspanスキーマを使用して、関数がList[ Document ]オブジェクトを返すことを確認します。
  • トレースを正しく構成するには、トレース名と retriever_schema 名が一致している必要があります。取得者スキーマの設定方法については、次のセクションを参照してください。

MLflowとの互換性を確認するために、リトリーバーのスキーマを設定します。​

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

取得者スキーマを手動で設定するには:

  1. コール 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"],
    )

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

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

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

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

Unity Catalogボリュームからファイルを読み込む​

エージェントがUnity Catalogボリュームに保存されている非構造化ファイル（テキスト文書、レポート、構成ファイルなど）を読み取る必要がある場合は、 Databricks SDK Files APIを使用してファイルを直接一覧表示および読み取るツールを作成できます。

以下の例では、エージェントが使用できる2つのツールを作成します。

  • list_volume_files : ボリューム内のファイルとディレクトリを一覧表示します。
  • read_volume_file : ボリュームからテキストファイルの内容を読み取ります。

Databricks AI Bridgeを含む最新バージョンのdatabricks-langchainをインストールしてください。

Bash
%pip install --upgrade databricks-langchain
Python
from databricks.sdk import WorkspaceClient
from langchain_core.tools import tool

VOLUME = "<catalog>.<schema>.<volume>"  # TODO: Replace with your volume
w = WorkspaceClient()


@tool
def list_volume_files(directory: str = "") -> str:
    """Lists files and directories in the Unity Catalog volume.
    Provide a relative directory path, or leave empty to list the volume root."""
    base = f"/Volumes/{VOLUME.replace('.', '/')}"
    path = f"{base}/{directory.lstrip('/')}" if directory else base
    entries = []
    for f in w.files.list_directory_contents(path):
        kind = "dir" if f.is_directory else "file"
        size = f" ({f.file_size} bytes)" if not f.is_directory else ""
        entries.append(f"  [{kind}] {f.name}{size}")
    return "\n".join(entries) if entries else "No files found."


@tool
def read_volume_file(file_path: str) -> str:
    """Reads a text file from the Unity Catalog volume.
    Provide the path relative to the volume root, for example 'reports/q1_summary.txt'."""
    base = f"/Volumes/{VOLUME.replace('.', '/')}"
    full_path = f"{base}/{file_path.lstrip('/')}"
    resp = w.files.download(full_path)
    return resp.contents.read().decode("utf-8")

ツールをLLMにバインドし、ツール呼び出しループを実行します。

Python
from databricks_langchain import ChatDatabricks
from langchain_core.messages import HumanMessage, ToolMessage

llm = ChatDatabricks(endpoint="databricks-claude-sonnet-4-5")
llm_with_tools = llm.bind_tools([list_volume_files, read_volume_file])

messages = [HumanMessage(content="What files are in the volume? Can you read about_databricks.txt and summarize it in 2 sentences?")]
tool_map = {"list_volume_files": list_volume_files, "read_volume_file": read_volume_file}

for _ in range(5):  # max iterations
    response = llm_with_tools.invoke(messages)
    messages.append(response)
    if not response.tool_calls:
        break
    for tc in response.tool_calls:
        result = tool_map[tc["name"]].invoke(tc["args"])
        messages.append(ToolMessage(content=result, tool_call_id=tc["id"]))

print(response.content)