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

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

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エンドポイントとインデックスの作成を参照してください。

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サービスを使用してエージェントをサードパーティツールに接続するを参照してください。

次の例では、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_urisimilarity_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-langchaindatabricks-openai などの Databricks AI Bridge パッケージを使用します。これらのパッケージには、既存の Databricks リソースからレトリーバーを作成するための from_vector_searchfrom_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を構成し、カスタムエンべディングモデルとテキスト列を指定する必要があります。エンべディングを提供するためのオプションを参照してください。

次の例は、VectorSearchRetrieverToolcolumnsembedding キーで設定する方法を示しています。

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.
)

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

ローカルツールが準備できたら、エージェントコードの一部として直接本番運用するか、検出可能性とガバナンスが向上しますが特定の制限があるUnity Catalog関数に移行できます。

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

UC関数を使用したDatabricks AI検索のクエリ(非推奨)

注記

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

Unity Catalog 関数を作成して、AI Search index クエリをラップできます。このアプローチ:

  • ガバナンスと検出可能性を備えた本番運用のユースケースをサポートします。
  • vector_search() SQL 関数を内部で使用します
  • 自動 MLflow トレースをサポートしています。
    • page_contentmetadataのエイリアスを使用して、関数の出力を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 AI 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トレース機能を追加して、リトリーバーを監視およびデバッグしましょう。トレースにより、実行の各ステップの入力、出力、およびメタデータを表示できます。

前の例では、@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_keytext_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)