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

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

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

Databricksは、AI Searchインデックスと外部ベクターストアの非構造化データにエージェントを接続するための複数のアプローチを提供します。Databricks AI Searchインデックスへの即時アクセスには事前構成済みのMCPサーバーを使用するか、AI Bridgeパッケージを使用してローカルでレトリーバーツールを開発するか、特殊なワークフロー向けにカスタムレトリーバー関数を構築します。

Databricks AI Searchは、以前はDatabricks ベクトル検索として知られていました。以前の/api/2.0/mcp/vector-search/ URLプレフィックスは、下位互換性のために引き続き機能します。

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

Databricks マネージド AI Search MCP サーバーを使用して、エージェントに Databricks AI Search インデックスへのアクセスを許可します。まず、Databricks マネージド埋め込みを使用してインデックスを作成します。AI Search Endpoint とインデックスの作成を参照してください。

AI Search の管理対象 MCP URL は https://<workspace-hostname>/api/2.0/mcp/ai-search/{catalog}/{schema}/{index_name} です。それに接続し、公開するツールをリストします：

Python
from databricks.sdk import WorkspaceClient
from databricks_mcp import DatabricksMCPClient

workspace_client = WorkspaceClient()
host = workspace_client.config.host

mcp_client = DatabricksMCPClient(
    server_url=f"{host}/api/2.0/mcp/ai-search/<catalog>/<schema>/<index-name>",
    workspace_client=workspace_client,
)
tools = mcp_client.list_tools()

このサーバーを使用するエージェントを構築してデプロイするには、エージェントで MCP サーバーを使用するを参照してください。インデックスの Unity Catalog 保護可能オブジェクトに対して、エージェントに SELECT を付与します。

その他のアプローチ​

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

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

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

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

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

    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 AI 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(),
     &quot;Content-Type&quot;: &quot;application/json&quot;,
   },
      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 AI Search レトリーバーツールをローカルで構築するには、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={&quot;text_column LIKE&quot;: &quot;Databricks&quot;}, # 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 関数に移行したりできますが、一定の制限があります。

UC 関数を使用して Databricks AI Search をクエリー (非推奨)

UC関数を使用したDatabricks AI Searchのクエリ（非推奨）​

注記

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

AI 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 AI Search index name here
    index => 'catalog.schema.databricks_docs_index',
    query => query,
    num_results => 5
  )

このレトリーバーツールを AI エージェントで使用するには、UCFunctionToolkit でラップします。MLflow Logs に 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 トレースを追加して、レトリーバーを監視およびデバッグしてください。トレースでは、実行の各ステップの入力、出力、メタデータを表示できます。

前の例では、@mlflow.trace デコレータを __call__ および解析メソッドの両方に追加します。デコレータは、関数が呼び出されたときに開始され、関数が戻るときに終了する スパン を作成します。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]:
    ...

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

  • MLflowレトリーバースパンスキーマを使用し、関数が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)
このページの見出し