エージェントを非構造化データに接続する
エージェントは、質問に回答し、コンテキストに応じた応答を提供するために、ドキュメントコレクション、ナレッジベース、またはテキストコーパスのような非構造化データを多くの場合クエリーする必要があります。
新しい Python エージェントの場合は、AI Bridge(推奨)を使用して AI Search インデックスを直接クエリーします。既存の MCP 統合については、レガシー ワークスペース Endpointを参照してください。外部ベクター ストアへの接続や、Unity Catalog ボリュームからのファイルの読み取りも可能です。
AI Bridge を使用したレトリーバーの開発(推奨)
Databricks AI Search インデックスをクエリーする新しい Python エージェントには、Databricks では AI Bridge をお勧めします。databricks-langchain または databricks-openai の VectorSearchRetrieverTool を使用して、インデックスをエージェント ツールとして公開します。
これらの例を実行する前に、ワークスペースにサインインし、クエリーできるAI Searchインデックスを選択します。例のインデックス名を 3 パートの名前で置き換えます。最初の例では、Databricksが管理するエンべディングと、ツール呼び出しをサポートするモデルEndpointを使用します。
ターミナルでインストールコマンドをランします。ノートブックでは、pipの代わりに%pipを使用し、インストール後にPythonを再起動します。
- LangChain/LangGraph
- OpenAI
databricks-langchain および LangGraph の最新バージョンをインストールします。
pip install --upgrade databricks-langchain langgraph
レトリーバーツールを作成し、検索のタイミングをエージェントに指示する説明を追加します。この例では、取得を直接テストしてから、ツールを呼び出してその結果を使用できるエージェントをランします。
from databricks_langchain import VectorSearchRetrieverTool, ChatDatabricks
from langgraph.prebuilt import create_react_agent
vs_tool = VectorSearchRetrieverTool(
index_name="catalog.schema.my_databricks_docs_index",
tool_name="databricks_docs_retriever",
tool_description="Search Databricks documentation to answer product questions.",
)
print(vs_tool.invoke("What is Databricks Agent Framework?"))
agent = create_react_agent(
ChatDatabricks(endpoint="databricks-claude-sonnet-4-5"),
tools=[vs_tool],
)
result = agent.invoke({"messages": [{
"role": "user",
"content": "Search the documentation and explain Databricks Agent Framework.",
}]})
for message in result["messages"]:
print(message)
印刷された会話でツール呼び出し、取得された文書、およびそれらの文書に基づく回答を確認します。
直接アクセス インデックスまたは自己管理型エンべディングを使用する Delta Sync インデックスを使用するシナリオでは、VectorSearchRetrieverTool を構成し、カスタム エンべディング モデルとテキスト列を指定する必要があります。埋め込みを提供するオプションを参照してください。
次の例は、columnsとembeddingキーでVectorSearchRetrieverToolを構成する方法を示しています。
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="databricks_docs_retriever", # 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 ドキュメントを参照してください。
databricks-openaiの最新バージョンをインストールします。
pip install --upgrade databricks-openai
この例では、OpenAI互換のDatabricksクライアントを使用します。モデルのツール呼び出しを実行し、取得したドキュメントをモデルに返します。
import json
from databricks_openai import DatabricksOpenAI, VectorSearchRetrieverTool
client = DatabricksOpenAI()
dbvs_tool = VectorSearchRetrieverTool(
index_name="catalog.schema.my_databricks_docs_index",
tool_name="databricks_docs_retriever",
tool_description="Search Databricks documentation to answer product questions.",
)
messages = [{
"role": "user",
"content": "Search the documentation and explain Databricks Agent Framework.",
}]
for _ in range(10):
response = client.chat.completions.create(
model="databricks-claude-sonnet-4-5",
messages=messages,
tools=[dbvs_tool.tool],
)
message = response.choices[0].message
messages.append(message.model_dump(exclude_none=True))
if not message.tool_calls:
print(message.content)
break
for call in message.tool_calls:
if call.function.name != dbvs_tool.tool["function"]["name"]:
raise ValueError(f"Unknown tool: {call.function.name}")
args = json.loads(call.function.arguments)
documents = dbvs_tool.execute(**args, openai_client=client)
print(documents)
messages.append({
"role": "tool",
"tool_call_id": call.id,
"content": json.dumps(documents),
})
else:
raise RuntimeError("The agent reached the tool-calling limit.")
直接アクセス インデックスまたは自己管理型エンべディングを使用する Delta Sync インデックスを使用するシナリオでは、VectorSearchRetrieverTool を構成し、カスタム エンべディング モデルとテキスト列を指定する必要があります。埋め込みを提供するオプションを参照してください。
次の例は、columnsとembeddingキーでVectorSearchRetrieverToolを構成する方法を示しています。
from databricks_openai import VectorSearchRetrieverTool
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="databricks_docs_retriever", # 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_model_name="databricks-bge-large-en" # The embedding model. Required for direct-access index or delta-sync index with self-managed embeddings.
)
Databricksの埋め込みEndpointを使用する場合は、上記のエイージェントループのように、openai_client=clientをvs_tool.execute()に渡します。インデックスベクトルを作成したのと同じ埋め込みモデルを使用します。
詳細については、VectorSearchRetrieverTool の API ドキュメントを参照してください。
リトリーバーをテストした後、それをカスタムエージェントに追加します。
Databricksの外部にあるベクトル検索インデックスをクエリーする
ベクトルインデックスが Databricks の外部でホストされている場合は、 Unity Catalog 接続を作成して外部サービスに接続し、エージェントコードでその接続を使用できます。外部 HTTP サービスへの接続を参照してください。
次の例では、HTTP 接続を介して Azure AI Search インデックスをクエリーします。そのインデックスをクエリーする権限を持つ接続と、接続に対する USE CONNECTION が必要です。Python の例を実行する前に、requests、databricks-sdk、および mlflow をインストールします。
-
HTTP接続のセットアップを使用して、Azure AI SearchへのUnity Catalog接続を作成します。この例ではベアラートークンを使用します。Azure AI Searchのデプロイに必要な認証方法を選択してください。
SQLCREATE CONNECTION azure_search_connection
TYPE HTTP
OPTIONS (
host 'https://example.search.windows.net',
base_path '/',
bearer_token secret ('<secret-scope>','<secret-key>')
); -
Unity Catalog接続を使用して、エージェントコードで取得ツールを定義します。この例では、MLflowデコレーターを使用してエージェントのトレースを有効にします。
MLflow レトリーバー スキーマに準拠するには、レトリーバー関数は List[Document] オブジェクトを返し、Document クラスの metadata フィールドを使用して、返されたドキュメントに doc_uri や similarity_score などの属性を追加する必要があります。MLflow ドキュメントをご覧ください。
import mlflow
import requests
from databricks.sdk import WorkspaceClient
from mlflow.entities import Document
class VectorSearchRetriever:
def __init__(self, connection_name, index_name):
self.workspace = WorkspaceClient()
self.connection_name = connection_name
self.index_name = index_name
@mlflow.trace(span_type="RETRIEVER", name="vector_search")
def __call__(self, query_vector: list[float], score_threshold: float = 0.0) -> list[Document]:
response = requests.post(
f"{self.workspace.config.host}/api/2.0/unity-catalog/connections/"
f"{self.connection_name}/proxy/indexes/{self.index_name}/docs/search",
params={"api-version": "2024-07-01"},
headers=self.workspace.config.authenticate(),
json={
"select": "HotelId,HotelName,Description,Category",
"vectorQueries": [{
"vector": query_vector,
"k": 7,
"fields": "DescriptionVector",
"kind": "vector",
"exhaustive": True,
}],
},
timeout=30,
)
response.raise_for_status()
return self.parse_results(response.json(), score_threshold)
@mlflow.trace(span_type="PARSER")
def parse_results(self, results, score_threshold) -> list[Document]:
return [
Document(
page_content=item.get("Description", ""),
id=item.get("HotelId"),
metadata={
"similarity_score": item.get("@search.score", 0),
"HotelName": item.get("HotelName"),
"Category": item.get("Category"),
},
)
for item in results.get("value", [])
if item.get("@search.score", 0) >= score_threshold
]
-
インデックスで使用されているものと同じモデルとディメンションを使用して、クエリーのエンベディングを生成します。数値配列を
query-vector.jsonとして保存し、レトリーバーのランを実行します:Pythonimport json
from pathlib import Path
query_vector = json.loads(Path("query-vector.json").read_text())
retriever = VectorSearchRetriever("azure_search_connection", "hotels_vector_index")
results = retriever(query_vector, score_threshold=0.1)
for document in results:
print(document.page_content)
レトリーバーツールにトレースを追加する
MLflow トレースを追加して、レトリーバーを監視およびデバッグしてください。トレースでは、実行の各ステップの入力、出力、メタデータを表示できます。
前の例では、@mlflow.trace デコレータを __call__ および解析メソッドの両方に追加します。デコレータは、関数が呼び出されたときに開始され、関数が戻るときに終了する スパン を作成します。MLflow は、関数の入力と出力、および発生した例外を自動的に記録します。
LangChain、LlamaIndex、OpenAI ライブラリのユーザーは、デコレーターを使用してトレースを手動で定義することに加えて、MLflow の自動ロギングを使用できます。トレーシングの概要を参照してください。
上記の例のように、page_content と metadata を持つ Document オブジェクトのリストを返します。これは MLflow レトリーバー スパン スキーマに準拠しているため、評価およびレビュー ツールで取得したドキュメントを認識できます。
Unity Catalogボリュームからファイルを読み込む
エージェントがUnity Catalogボリュームに保存されている非構造化ファイル(テキストドキュメント、レポート、設定ファイルなど)を読み取る必要がある場合は、Databricks SDKのFiles APIを使用してファイルを直接一覧表示および読み取るツールを作成できます。
既存のボリュームを選択し、エージェントのアイデンティティに USE CATALOG、USE SCHEMA、および READ VOLUME を付与します。コード内の <catalog>.<schema>.<volume> を置き換えます。サンプルタスクについては、about_databricks.txt という名前の UTF-8 テキストファイルを使用するか、プロンプトをボリューム内のファイルに変更します。
次の例では、エージェントが使用できる2つのツールを作成します:
list_volume_files:ボリューム内のファイルとディレクトリを一覧表示します。read_volume_file:ボリュームからテキストファイルの内容を読み取ります。
- LangChain/LangGraph
- OpenAI
Databricks AI Bridgeを含むdatabricks-langchainの最新バージョンをインストールします。
pip install --upgrade databricks-langchain langgraph
from pathlib import PurePosixPath
from databricks.sdk import WorkspaceClient
from langchain_core.tools import tool
VOLUME = "<catalog>.<schema>.<volume>"
w = WorkspaceClient()
def volume_path(relative_path: str = "") -> str:
path = PurePosixPath(relative_path)
if path.is_absolute() or ".." in path.parts:
raise ValueError("Use a path relative to the configured volume.")
return str(PurePosixPath("/Volumes", *VOLUME.split(".")) / path)
@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."""
path = volume_path(directory)
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'."""
full_path = volume_path(file_path)
resp = w.files.download(full_path)
return resp.contents.read().decode("utf-8")
ツールを LLM にバインドし、ツール呼び出しループを実行します:
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"]))
else:
raise RuntimeError("The agent reached the tool-calling limit.")
print(response.content)
Databricks AI Bridgeを含むdatabricks-openaiの最新バージョンをインストールします。
pip install --upgrade databricks-openai
import json
from pathlib import PurePosixPath
from databricks.sdk import WorkspaceClient
from databricks_openai import DatabricksOpenAI
VOLUME = "<catalog>.<schema>.<volume>"
w = WorkspaceClient()
client = DatabricksOpenAI()
# Define the tool specifications
tools = [
{
"type": "function",
"function": {
"name": "list_volume_files",
"description": "Lists files and directories in the Unity Catalog volume. Provide a relative directory path, or leave empty to list the volume root.",
"parameters": {
"type": "object",
"properties": {
"directory": {
"type": "string",
"description": "Relative directory path within the volume. Leave empty for root.",
}
},
"required": [],
},
},
},
{
"type": "function",
"function": {
"name": "read_volume_file",
"description": "Reads a text file from the Unity Catalog volume. Provide the path relative to the volume root, for example 'reports/q4_summary.txt'.",
"parameters": {
"type": "object",
"properties": {
"file_path": {
"type": "string",
"description": "Path to the file relative to the volume root.",
}
},
"required": ["file_path"],
},
},
},
]
def volume_path(relative_path: str = "") -> str:
path = PurePosixPath(relative_path)
if path.is_absolute() or ".." in path.parts:
raise ValueError("Use a path relative to the configured volume.")
return str(PurePosixPath("/Volumes", *VOLUME.split(".")) / path)
def execute_tool(name: str, args: dict) -> str:
if name == "list_volume_files":
directory = args.get("directory", "")
path = volume_path(directory)
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."
elif name == "read_volume_file":
full_path = volume_path(args["file_path"])
resp = w.files.download(full_path)
return resp.contents.read().decode("utf-8")
return f"Unknown tool: {name}"
# Call the model with tools
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "List the files in the volume, then read about_databricks.txt and summarize it."},
]
for _ in range(10):
response = client.chat.completions.create(
model="databricks-claude-sonnet-4-5", messages=messages, tools=tools
)
message = response.choices[0].message
messages.append(message.model_dump(exclude_none=True))
if not message.tool_calls:
print(message.content)
break
for call in message.tool_calls:
args = json.loads(call.function.arguments)
result = execute_tool(call.function.name, args)
messages.append({"role": "tool", "tool_call_id": call.id, "content": result})
else:
raise RuntimeError("The agent reached the tool-calling limit.")
MCP(レガシー)を使用してAI Searchインデックスをクエリーする
AI Search MCP サーバーは、レガシー ワークスペース MCP Endpointを使用します。Databricks 管理埋め込みを持つインデックスが必要です。AI Search Endpointとインデックスの作成を参照してください。
AI SearchのレガシーワークスペースMCP URLは https://<workspace-hostname>/api/2.0/mcp/ai-search/{catalog}/{schema}/{index_name} です。接続して、公開されているツールを一覧表示します。
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()
このサーバーを使用するエージェントを構築してデプロイするには、PythonエージェントでのMCPツールの使用を参照してください。インデックスの Unity Catalog 保護対象に対する SELECT をエージェントに付与します。
AI Search は以前はベクトル検索でした。以前の /api/2.0/mcp/vector-search/ URL プレフィックスも引き続き機能します。
UC関数を使用したDatabricks AI Searchのクエリ(非推奨)
新しいPythonエージェントには、AI Bridgeを使用します。次の例は、既存の統合の参照用です。
AI Search インデックスのクエリーをラップする Unity Catalog 関数を作成できます。このアプローチ:
- ガバナンスと検出可能性を備えた本番運用ユースケースをサポートします
- vector_search() SQL関数を内部で使用します。
- 自動MLflow トレースをサポートしています。
- 関数の出力をMLflow レトリーバースキーマに揃えるには、
page_contentおよびmetadataのエイリアスを使用する必要があります。 - 最上位の出力キーとしてではなく、「SQL マップ関数」を使用して、追加のメタデータ列を
metadata列に追加する必要があります。
- 関数の出力をMLflow レトリーバースキーマに揃えるには、
関数を作成するには、ノートブックまたは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
)
関数ツールのセットアップに従い、このレトリーバーを UCFunctionToolkit でラップしてエージェントで使用します。これにより、MLflow の Logs 内で RETRIEVER スパンタイプが自動的に生成され、MLflow を通じた自動トレースが可能になります。
from databricks_langchain import UCFunctionToolkit
from unitycatalog.ai.core.databricks import DatabricksFunctionClient
toolkit = UCFunctionToolkit(
function_names=[
"main.default.databricks_docs_vector_search"
],
client=DatabricksFunctionClient(),
)
tools = toolkit.tools
Unity Catalogレトリーバーツールには、次の注意点があります:
- SQLクライアントは、返される行またはバイトの最大数を制限する場合があります。データの切り捨てを防ぐには、UDFによって返される列の値を切り捨てます。たとえば、
substring(chunked_text, 0, 8192)を使用することで、大きなコンテンツ列のサイズを縮小し、実行中の行の切り捨てを回避できます。 - このツールは
vector_search()関数のラッパーであるため、vector_search()関数と同じ制限を受けます。制限事項を参照してください。
UCFunctionToolkitの詳細については、Unity Catalog のドキュメントを参照してください。
カスタム リトリーバ スキーマ (レガシー)
mlflow.models.set_retriever_schema は非推奨です。新しいレトリーバーの場合は、標準レトリーバースキーマを使用して Document オブジェクトを返します。次の構成は、カスタムフィールドを返す既存の統合用です。
リトリーバーのスキーマを手動で設定するには:
-
mlflow.models.set_retriever_schema を呼び出します。エージェントを定義する場合。
set_retriever_schemaを使用して、返されたテーブルの列名を MLflow の想定フィールド (primary_key、text_column、doc_uriなど) にマッピングします。Pythonimport mlflow
# 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を使用して、リトリーバーの結果を事前に定義された評価データセットと比較し、リトリーバーの再現率と精度を決定します。Evaluation データセット reference を参照してください。