本書は著者が手動で翻訳したものであり内容の正確性を保証するものではありません。正確な内容に関しては原文を参照ください。
著者: Bo cheng, Solutions Architect and Charlie Hohenstein, Solutions Architect
イントロダクション
AIを活用して、「地域ごとの売上を棒グラフで見せて」といったシンプルな質問をデータに尋ねて、即座にデータとインタラクティブな可視化を受け取るといったことを想像してください。SQLは不要です。手動のチャートの設定も不要です。単に会話だけです。
この記事では、以下を用いたデータチャットbotの構築方法を説明します:
- Databricks Genie — 自然言語のデータクエリーのためのマネージドMCP
- Unity Catalogの関数 — AIが呼び出せるツールとしてのマネージドMCP
- Claude 3.7 SonnetとDatabricksモデルサービング
- 堅牢なエージェントデプロイメントのためのMosaic AI Agent FrameworkとLangGraph
- フロントエンドUIのためのDatabricks Apps
最後までには、あなたのデータウェアハウスにクエリーし、チャートを生成し、結果を表示できる完全に動作するチャットbotを手に入れることになります。
課題
データチームは共通の課題に直面します: ビジネスユーザーはクイックな洞察を欲しますが、データへの問い合わせにはSQLの知識を必要とし、可視化の作成にはさらに技術的なスキルを必要とします。従来のBIツールは助けになりますが、依然として学習曲線が存在します。シンプルな英語でユーザーが質問できるとしたらどうでしょうか？
ここで、AIエージェントの登場です。しかし、以下を可能とするエージェントをどのようにすれば構築できるのでしょうか:
- 自然言語の問い合わせを理解する
- あなたのデータウェアハウスに対してSQLを実行する
- 結果を可視化に変換する
- 会話のコンテキストを保持する
- エージェントのパフォーマンスや品質を監視する
ソリューション
このAIエージェントシステムを構成するいくつかの主要コンポーネントがあります:
- チャットモデル(LLM): チャットモデルは入力としてメッセージのリストを受け取り、出力としてメッセージを返却するインタフェースを提供します。LLMの最新のイテレーションでは、ツール呼び出しやMCPサーバーによるやり取りが可能になっています。
- フレームワーク: Langgraphのような低レベルのオーケストレーションフレームワークと、databricks-agentsのようなデプロイメントフレームワークを連携させることで、LLMとさまざまな(外部、内部)MCPサーバーの間のインタラクションを容易に構築し、抽象化することができます。
- マネージドMCP: Model Context Protocol (MCP)サーバーはブリッジとして動作し、LLMがバイブのデータやツールにアクセスできるようにします。この記事では、構造化データおよび最終的には当該データのプロットを返却するために、Genieスペースを活用し、UCの関数が管理するDatabricks MCPサーバーを活用します。
Databricks Apps: Databricksプラットフォーム上に直接デプロイされ、Unity Catalog、DBSQL、Mosaic AIモデルサービングとネイティブに連携するAIアプリです - エンドユーザーがエンドポイントとチャットし、我々のプロットツールのJSON表現を可視化することを可能にするエージェントエンドポイントです。
以下のデモでは、テキストから可視化に変換する能力を提供するために、我々がどのようにMosaic AIエージェントフレームワークとUCの関数、GenieスペースのDatabricksマネージドMCPサーバーを組み合わせたのかを説明します。Mosaic AIモデルsーあビングとDatabricks-agentsのデプロイメントフレームワークを用いてモデルをサービングする能力は、アプリやユーザー認証を用いてエンドユーザーにサービスを提供するDatabricks Appに対する簡単なインタフェースを提供します。
注意: このデモにおいては以下のライブラリを使用しています。セットアップの観点では、ノートブックにおけるDatabricksサーバレスコンピュートを活用しています。記事の最後には、カバーされているすべてのコードのGitHubリポジトリを含めています。
databricks-langchain
databricks-mcp
langgraph==0.5.3
uv
databricks-agents
mlflow-skinny[databricks]
unitycatalog-ai[databricks]
マネージドMCPサーバー
自然言語の質問を用いて構造化データテーブルから洞察を得るためにGenieスペースにクエリーを行います
我々の例では、エージェントが(マネージドMCP経由で)samples.nyctaxi.tripsテーブルからデータを選択する既存のGenieスペースを呼び出せるようにしています。
そして、我々のレスポンスを生成するUC関数を(マネージドMCP経由で)呼び出すために、Genieスペースからのレスポンスに含まれるデータを活用します:
- genie_to_chart: PythonのUC関数は、Databricksランタイムバージョン17以降のDBSQLサーバレスウェアハウスにおいて、関数実行環境で依存関係を指定できるようになりました。入力としてのGenieのレスポンスを出力としてのjson文字列として表現されるplotlyの図に変換するために、依存関係としてplotlyとpandasを活用します。
これらをDBSQLのコードとして、Databricks SQLエディタで直接実行することができ、あるいはUC関数を作成するためにDatabricksノートブックのSQLセルを使うことができます:
from unitycatalog.ai.core.databricks import DatabricksFunctionClient
client = DatabricksFunctionClient()
CATALOG = dbutils.widgets.get("catalog")
SCHEMA = dbutils.widgets.get("schema")
SPACE_ID = dbutils.widgets.get("space_id")
def genie_to_chart(genie_response_json: str, chart_type: str) -> str:
"""
Transform Genie MCP response into a Plotly chart.
This function handles all data extraction and transformation from Genie's nested JSON format.
Args:
genie_response_json (str): Raw JSON string from Genie MCP query_space tool
chart_type (str): Type of chart - "bar", "line", or "pie"
Returns:
str: Plotly JSON string for direct rendering with Plotly
"""
import json
import plotly.express as px
import pandas as pd
# Parse Genie response - handle both full wrapper and extracted content
genie_data = json.loads(genie_response_json)
# If there's a 'content' field, parse it (double-parse scenario)
if "content" in genie_data:
content_data = (
json.loads(genie_data["content"])
if isinstance(genie_data["content"], str)
else genie_data["content"]
)
else:
# Already the content (no wrapper)
content_data = genie_data
# Extract columns and rows from statement_response
columns_info = content_data["statement_response"]["manifest"]["schema"]["columns"]
column_names = [col["name"] for col in columns_info]
data_array = content_data["statement_response"]["result"]["data_array"]
rows = [[value["string_value"] for value in row["values"]] for row in data_array]
# Create DataFrame
df = pd.DataFrame(rows, columns=column_names)
# Convert all columns to appropriate types
for col in df.columns:
df[col] = pd.to_numeric(df[col], errors="ignore")
# Column selection: string columns for X, last numeric for Y
string_cols = df.select_dtypes(include=["object", "string"]).columns
numeric_cols = df.select_dtypes(include="number").columns
x_col = string_cols[0] if len(string_cols) > 0 else df.columns[0]
y_col = numeric_cols[-1] if len(numeric_cols) > 0 else df.columns[-1]
# Generate chart
chart_functions = {
"bar": lambda: px.bar(df, x=x_col, y=y_col, title=f"{y_col} by {x_col}"),
"line": lambda: px.line(
df, x=x_col, y=y_col, title=f"{y_col} by {x_col}", markers=True
),
"pie": lambda: px.pie(
df, names=x_col, values=y_col, title=f"{y_col} by {x_col}"
),
}
chart_type_key = chart_type.lower().strip()
if chart_type_key not in chart_functions:
raise ValueError(f"Unsupported chart type: {chart_type}")
fig = chart_functions[chart_type_key]()
# Return structured response with Plotly JSON
response = {
"plotly_json": json.loads(fig.to_json()),
"chart_type": chart_type_key,
}
return json.dumps(response)
# Create the function and supply the dependency in standard PyPI format
client.create_python_function(
func=genie_to_chart,
catalog=CATALOG,
schema=SCHEMA,
replace=True,
dependencies=["plotly", "pandas"],
)
関数を作成したら、指定したカタログとスキーマの関数タブ配下でそれらを確認できるはずです。
GenieスペースとUC関数を作成すると、Databricksは自動的にこれらをMCPサーバー経由で利用できるようにします。なので、MCPサーバーを使用するエージェントを作成する必要があります。
Mosaic AI Agent Framework
このソリューションでは、databricks-mcpライブラリのような構築済みのインテげうレーションを用いて、LangGraphエージェントに引き渡し可能なツールとして、DatabricksマネージドMCPサーバーをインテグレーションする方法をカバーします。
コードとしてのエージェントモデルにおいて、以下を行う柔軟なツール呼び出しエージェントを作成します:
- MCPツールをラッピングする
- エージェントがマネージドMCPサーバーで利用可能なツールを自動的に発見可能にする
- ResponsesAgentクラスを用いてLangGraphエージェントを定義する
- MLflowオートトレーシングとインテグレーションする
Databricks Appsに対するモデルエンドポイントとして我々のエージェントを公開するためには、コードで開発される我々のエージェントを容易にデプロイするためにMosaic AI Agent Frameworkを活用します。エージェントのモデルサービングエンドポイントをデプロイすると、テストするためにAI Playgroundを活用したり、専門家がフィードバックを提供するためにレビューアプリを活用できるようになります。
エージェントのモデルサービングエンドポイントをデプロイする前に、上述したMosaic AI Agent Frameworkやモデルサービングエンドポイントと完全な互換性を持つ、マルチターンの会話型エージェントをシンプルにするために、(以下のコードで開発しているエージェントで使用されている)MLflowのResponsesAgentクラスを使用することをお勧めします。
import asyncio
import mlflow
import os
import json
from uuid import uuid4
from pydantic import BaseModel, create_model
from typing import Annotated, Any, Generator, List, Optional, Sequence, TypedDict, Union
from databricks_langchain import (
ChatDatabricks,
UCFunctionToolkit,
VectorSearchRetrieverTool,
)
from databricks_mcp import DatabricksOAuthClientProvider, DatabricksMCPClient
from databricks.sdk import WorkspaceClient
from langchain_core.language_models import LanguageModelLike
from langchain_core.runnables import RunnableConfig, RunnableLambda
from langchain_core.messages import (
AIMessage,
AIMessageChunk,
BaseMessage,
convert_to_openai_messages,
)
from langchain_core.tools import BaseTool, tool
from langgraph.graph import END, StateGraph
from langgraph.graph.message import add_messages
from langgraph.graph.state import CompiledStateGraph
from langgraph.prebuilt.tool_node import ToolNode
from mcp import ClientSession
from mcp.client.streamable_http import streamablehttp_client as connect
from mlflow.entities import SpanType
from mlflow.pyfunc import ResponsesAgent
from mlflow.types.responses import (
ResponsesAgentRequest,
ResponsesAgentResponse,
ResponsesAgentStreamEvent,
)
import nest_asyncio
nest_asyncio.apply()
############################################
## Define your LLM endpoint and system prompt
############################################
# TODO: Replace with your model serving endpoint
LLM_ENDPOINT_NAME = "databricks-claude-3-7-sonnet"
llm = ChatDatabricks(endpoint=LLM_ENDPOINT_NAME)
# TODO: Update with your system prompt
system_prompt = """
You are a helpful assistant that can query data and create charts.
WORKFLOW:
1. If user asks for data or a chart, call query_space_01f0ab8c079d17b8a00584e70d2ac18c to get data
2. If user wants a chart, call chat_app_demo__dev__genie_to_chart EXACTLY ONCE with:
- genie_response_json: the ENTIRE raw JSON response from query_space
- chart_type: "bar", "line", or "pie"
3. IMMEDIATELY STOP after calling genie_to_chart and provide a brief summary
CRITICAL RULES:
- Use EXACT tool names: query_space_01f0ab8c079d17b8a00584e70d2ac18c and chat_app_demo__dev__genie_to_chart
- Call genie_to_chart ONLY ONCE per request - NEVER call it multiple times
- STOP calling tools after genie_to_chart returns - the chart is already created
- Do NOT query data again or recreate charts - ONE chart is sufficient"""
###############################################################################
## Configure MCP Servers for your agent
## This section sets up server connections so your agent can retrieve data or take actions.
###############################################################################
# TODO: Choose your MCP server connection type.
# ----- Simple: Managed MCP Server (no extra setup required) -----
# Uses your Databricks Workspace settings and Personal Access Token (PAT) auth.
workspace_client = WorkspaceClient()
CATALOG = "catalog"
SCHEMA = "schema"
GENIE_SPACE_ID = "space"
# Managed MCP Servers: Ready to use with default settings above
host = workspace_client.config.host
MANAGED_MCP_SERVER_URLS = [
f"{host}/api/2.0/mcp/functions/{CATALOG}/{SCHEMA}",
f"{host}/api/2.0/mcp/genie/{GENIE_SPACE_ID}",
]
# ----- Advanced (optional): Custom MCP Server with OAuth -----
# For Databricks Apps hosting custom MCP servers, OAuth with a service principal is required.
# Uncomment and fill in your settings ONLY if connecting to a custom MCP server.
#
# import os
# workspace_client = WorkspaceClient(
# host="",
# client_id=os.getenv("DATABRICKS_CLIENT_ID"),
# client_secret=os.getenv("DATABRICKS_CLIENT_SECRET"),
# auth_type="oauth-m2m", # Enables machine-to-machine OAuth
# )
# Custom MCP Servers: Add URLs below if needed (requires custom setup and OAuth above)
CUSTOM_MCP_SERVER_URLS = [
# Example: "https:///mcp"
]
#####################
## MCP Tool Creation
#####################
# Define a custom LangChain tool that wraps functionality for calling MCP servers
class MCPTool(BaseTool):
"""Custom LangChain tool that wraps MCP server functionality"""
def __init__(
self,
name: str,
description: str,
args_schema: type,
server_url: str,
ws: WorkspaceClient,
is_custom: bool = False,
):
# Initialize the tool
super().__init__(name=name, description=description, args_schema=args_schema)
# Store custom attributes: MCP server URL, Databricks workspace client, and whether the tool is for a custom server
object.__setattr__(self, "server_url", server_url)
object.__setattr__(self, "workspace_client", ws)
object.__setattr__(self, "is_custom", is_custom)
def _run(self, **kwargs) -> str:
"""Execute the MCP tool"""
if self.is_custom:
# Use the async method for custom MCP servers (OAuth required)
return asyncio.run(self._run_custom_async(**kwargs))
else:
# Use managed MCP server via synchronous call
mcp_client = DatabricksMCPClient(
server_url=self.server_url, workspace_client=self.workspace_client
)
response = mcp_client.call_tool(self.name, kwargs)
return "".join([c.text for c in response.content])
async def _run_custom_async(self, **kwargs) -> str:
"""Execute custom MCP tool asynchronously"""
async with connect(
self.server_url, auth=DatabricksOAuthClientProvider(self.workspace_client)
) as (
read_stream,
write_stream,
_,
):
# Create an async session with the server and call the tool
async with ClientSession(read_stream, write_stream) as session:
await session.initialize()
response = await session.call_tool(self.name, kwargs)
return "".join([c.text for c in response.content])
# Retrieve tool definitions from a custom MCP server (OAuth required)
async def get_custom_mcp_tools(ws: WorkspaceClient, server_url: str):
"""Get tools from a custom MCP server using OAuth"""
async with connect(server_url, auth=DatabricksOAuthClientProvider(ws)) as (
read_stream,
write_stream,
_,
):
async with ClientSession(read_stream, write_stream) as session:
await session.initialize()
tools_response = await session.list_tools()
return tools_response.tools
# Retrieve tool definitions from a managed MCP server
def get_managed_mcp_tools(ws: WorkspaceClient, server_url: str):
"""Get tools from a managed MCP server"""
mcp_client = DatabricksMCPClient(server_url=server_url, workspace_client=ws)
return mcp_client.list_tools()
# Convert an MCP tool definition into a LangChain-compatible tool
def create_langchain_tool_from_mcp(
mcp_tool, server_url: str, ws: WorkspaceClient, is_custom: bool = False
):
"""Create a LangChain tool from an MCP tool definition"""
schema = mcp_tool.inputSchema.copy()
properties = schema.get("properties", {})
required = schema.get("required", [])
# Map JSON schema types to Python types for input validation
TYPE_MAPPING = {"integer": int, "number": float, "boolean": bool}
field_definitions = {}
for field_name, field_info in properties.items():
field_type_str = field_info.get("type", "string")
field_type = TYPE_MAPPING.get(field_type_str, str)
if field_name in required:
field_definitions[field_name] = (field_type, ...)
else:
field_definitions[field_name] = (field_type, None)
# Dynamically create a Pydantic schema for the tool's input arguments
args_schema = create_model(f"{mcp_tool.name}Args", **field_definitions)
# Return a configured MCPTool instance
return MCPTool(
name=mcp_tool.name,
description=mcp_tool.description or f"Tool: {mcp_tool.name}",
args_schema=args_schema,
server_url=server_url,
ws=ws,
is_custom=is_custom,
)
# Gather all tools from managed and custom MCP servers into a single list
async def create_mcp_tools(
ws: WorkspaceClient,
managed_server_urls: List[str] = None,
custom_server_urls: List[str] = None,
) -> List[MCPTool]:
"""Create LangChain tools from both managed and custom MCP servers"""
tools = []
if managed_server_urls:
# Load managed MCP tools
for server_url in managed_server_urls:
try:
mcp_tools = get_managed_mcp_tools(ws, server_url)
for mcp_tool in mcp_tools:
tool = create_langchain_tool_from_mcp(
mcp_tool, server_url, ws, is_custom=False
)
tools.append(tool)
except Exception as e:
print(f"Error loading tools from managed server {server_url}: {e}")
if custom_server_urls:
# Load custom MCP tools (async)
for server_url in custom_server_urls:
try:
mcp_tools = await get_custom_mcp_tools(ws, server_url)
for mcp_tool in mcp_tools:
tool = create_langchain_tool_from_mcp(
mcp_tool, server_url, ws, is_custom=True
)
tools.append(tool)
except Exception as e:
print(f"Error loading tools from custom server {server_url}: {e}")
return tools
#####################
## Define agent logic
#####################
# The state for the agent workflow, including the conversation and any custom data
class AgentState(TypedDict):
messages: Annotated[Sequence[BaseMessage], add_messages]
custom_inputs: Optional[dict[str, Any]]
custom_outputs: Optional[dict[str, Any]]
# Define the LangGraph agent that can call tools
def create_tool_calling_agent(
model: LanguageModelLike,
tools: Union[ToolNode, Sequence[BaseTool]],
system_prompt: Optional[str] = None,
):
model = model.bind_tools(tools) # Bind tools to the model
# Function to check if agent should continue or finish based on last message
def should_continue(state: AgentState):
messages = state["messages"]
last_message = messages[-1]
# If function (tool) calls are present, continue; otherwise, end
if isinstance(last_message, AIMessage) and last_message.tool_calls:
return "continue"
else:
return "end"
# Preprocess: optionally prepend a system prompt to the conversation history
if system_prompt:
preprocessor = RunnableLambda(
lambda state: [{"role": "system", "content": system_prompt}]
+ state["messages"]
)
else:
preprocessor = RunnableLambda(lambda state: state["messages"])
model_runnable = preprocessor | model # Chain the preprocessor and the model
# The function to invoke the model within the workflow
def call_model(
state: AgentState,
config: RunnableConfig,
):
response = model_runnable.invoke(state, config)
return {"messages": [response]}
workflow = StateGraph(AgentState) # Create the agent's state machine
workflow.add_node("agent", RunnableLambda(call_model)) # Agent node (LLM)
workflow.add_node("tools", ToolNode(tools)) # Tools node
workflow.set_entry_point("agent") # Start at agent node
workflow.add_conditional_edges(
"agent",
should_continue,
{
"continue": "tools", # If the model requests a tool call, move to tools node
"end": END, # Otherwise, end the workflow
},
)
workflow.add_edge("tools", "agent") # After tools are called, return to agent node
# Compile and return the tool-calling agent workflow
return workflow.compile()
# ResponsesAgent class to wrap the compiled agent and make it compatible with Mosaic AI Responses API
class LangGraphResponsesAgent(ResponsesAgent):
def __init__(self, agent):
self.agent = agent
# Convert a LangChain message to a Responses-format dictionary
def _langchain_to_responses(
self, messages: list[BaseMessage]
) -> list[dict[str, Any]]:
"""Convert from LangChain messages to Responses output item dictionaries"""
for message in messages:
message = message.model_dump() # Convert the message model to dict
role = message["type"]
if role == "ai":
if tool_calls := message.get("tool_calls"):
# Return function call items for all tool calls present
return [
self.create_function_call_item(
id=message.get("id") or str(uuid4()),
call_id=tool_call["id"],
name=tool_call["name"],
arguments=json.dumps(tool_call["args"]),
)
for tool_call in tool_calls
]
else:
# Regular AI text message
return [
self.create_text_output_item(
text=message["content"],
id=message.get("id") or str(uuid4()),
)
]
elif role == "tool":
# Output from tool/function execution
return [
self.create_function_call_output_item(
call_id=message["tool_call_id"],
output=message["content"],
)
]
elif role == "user":
# User messages as-is
return [message]
# Make a prediction (single-step) for the agent
def predict(self, request: ResponsesAgentRequest) -> ResponsesAgentResponse:
outputs = [
event.item
for event in self.predict_stream(request)
if event.type == "response.output_item.done" or event.type == "error"
]
return ResponsesAgentResponse(
output=outputs, custom_outputs=request.custom_inputs
)
# Stream predictions for the agent, yielding output as it's generated
def predict_stream(
self,
request: ResponsesAgentRequest,
) -> Generator[ResponsesAgentStreamEvent, None, None]:
cc_msgs = self.prep_msgs_for_cc_llm([i.model_dump() for i in request.input])
# Stream events from the agent graph
for event in self.agent.stream(
{"messages": cc_msgs}, stream_mode=["updates", "messages"]
):
if event[0] == "updates":
# Stream updated messages from the workflow nodes
for node_data in event[1].values():
if "messages" in node_data:
for item in self._langchain_to_responses(node_data["messages"]):
yield ResponsesAgentStreamEvent(
type="response.output_item.done", item=item
)
elif event[0] == "messages":
# Stream generated text message chunks
try:
chunk = event[1][0]
if isinstance(chunk, AIMessageChunk) and (content := chunk.content):
yield ResponsesAgentStreamEvent(
**self.create_text_delta(delta=content, item_id=chunk.id),
)
except:
pass
# Initialize the entire agent, including MCP tools and workflow
def initialize_agent():
"""Initialize the agent with MCP tools"""
# Create MCP tools from the configured servers
mcp_tools = asyncio.run(
create_mcp_tools(
ws=workspace_client,
managed_server_urls=MANAGED_MCP_SERVER_URLS,
custom_server_urls=CUSTOM_MCP_SERVER_URLS,
)
)
# Create the agent graph with an LLM, tool set, and system prompt (if given)
agent = create_tool_calling_agent(llm, mcp_tools, system_prompt)
return LangGraphResponsesAgent(agent)
mlflow.langchain.autolog()
AGENT = initialize_agent()
mlflow.models.set_model(AGENT)
そして、MLflowにおけるmodels from codeアプローチを用いて、エージェントを記録、登録する必要があります。
import mlflow
from agent import LLM_ENDPOINT_NAME
from mlflow.models.resources import (
DatabricksServingEndpoint,
DatabricksFunction,
DatabricksGenieSpace,
DatabricksTable,
)
from mlflow.models.auth_policy import AuthPolicy, SystemAuthPolicy, UserAuthPolicy
from pkg_resources import get_distribution
resources = [
DatabricksServingEndpoint(endpoint_name=LLM_ENDPOINT_NAME),
DatabricksFunction(function_name=f"{CATALOG}.{SCHEMA}.genie_to_chart"),
DatabricksGenieSpace(genie_space_id=SPACE_ID),
DatabricksTable(table_name="samples.nyctaxi.trips"),
]
# System policy: resources accessed with system credentials
system_policy = SystemAuthPolicy(resources=resources)
# User policy: API scopes for OBO access
api_scopes = [
"sql.statement-execution",
"mcp.genie",
"mcp.external",
"catalog.connections",
"mcp.vectorsearch",
"vectorsearch.vector-search-indexes",
"iam.current-user:read",
"sql.warehouses",
"dashboards.genie",
"serving.serving-endpoints",
"iam.access-control:read",
"apps.apps",
"mcp.functions",
"vectorsearch.vector-search-endpoints",
]
user_policy = UserAuthPolicy(api_scopes=api_scopes)
input_example = {
"input": [
{
"role": "user",
"content": "How many trips were taken each day in 2016? Show me a line chart",
}
]
}
with mlflow.start_run():
logged_agent_info = mlflow.pyfunc.log_model(
name="agent",
python_model="agent.py",
input_example=input_example,
resources=resources,
pip_requirements=[
"databricks-mcp",
f"langgraph=={get_distribution('langgraph').version}",
f"mcp=={get_distribution('mcp').version}",
f"databricks-langchain=={get_distribution('databricks-langchain').version}",
],
# auth_policy=AuthPolicy(
# system_auth_policy=system_policy, user_auth_policy=user_policy
# ),
)
mlflow.set_registry_uri("databricks-uc")
# TODO: define the catalog, schema, and model name for your UC model
model_name = "managed-mcp-model"
UC_MODEL_NAME = f"{CATALOG}.{SCHEMA}.{model_name}"
# register the model to UC
uc_registered_model_info = mlflow.register_model(
model_uri=logged_agent_info.model_uri, name=UC_MODEL_NAME
)
最後に、モデルサービングエンドポイントを作成し、専門家の評価のためのレビューアプリを作成するために、Mosaic AI Agent Frameworkのagents.deploy() を使用することができます。ここでは、Databricksホスト、クライアント、クライアントのシークレット(これは認証のためにサービスプリンシパルが使用しています)のようにエージェントが認証に必要とするすべての環境変数を記録しています。
from databricks import agents
agents.deploy(
UC_MODEL_NAME,
uc_registered_model_info.version,
# ==============================================================================
# TODO: ONLY UNCOMMENT AND CONFIGURE THE ENVIRONMENT_VARS SECTION BELOW
# IF YOU ARE USING OAUTH/SERVICE PRINCIPAL FOR CUSTOM MCP SERVERS.
# For managed MCP (the default), LEAVE THIS SECTION COMMENTED OUT.
# ==============================================================================
environment_vars={
"DATABRICKS_HOST": "{{secrets/dbdemos/DATABRICKS_HOST}}",
"DATABRICKS_CLIENT_ID": "{{secrets/dbdemos/DATABRICKS_CLIENT_ID}}",
"DATABRICKS_CLIENT_SECRET": "{{secrets/dbdemos/DATABRICKS_CLIENT_SECRET}}",
},
tags={"endpointSource": "docs"},
)
エージェントモデルエンドポイントの準備ができるまで待つ(約10-20分)と、専門家フィードバックのためのレビューアプリやさまざまなクエリーでテストを開始するためのAI Playgroundを利用できるようになります。AI Playgroundでは、エンドポイントの選択でカスタムエージェントを選択する必要があります。
AI Playgroundにおけるエージェントエンドポイントの使い方の手順はこちらです。
レビューアプリやAI Playgroundはフィードバックやテストでは非常に有用ですが、この例ではモデルサービングエンドポイントがDatabricks Appと連携させる方法を説明したいと思います。
Databricks Apps
注意: このDatabricks Appのコードはリポジトリに格納されています。以下ではメインロジックのスニペットを示しています。
チャットbotインタフェースを作成するために、フロントエンドアプリケーションをホスティングするためにDatabricks Appsを活用します。
Streamlitを用いてこのUIを作成し、DatabricksとMLflow Deployments SDKを用いてエージェントエンドポイントと連携させます。以下の機能を実行できるUIを必要としています:
- Databricks認証を用いてエージェントエンドポイントを呼び出し
- エージェントのレスポンスを解析
- エージェントのレスポンスの表示、可視化のレンダリング
Databricks認証によるエージェントエンドポイントの呼び出し
エージェントに接続して質問を行うためには、エージェントクライアントを作成する必要があります。エージェントクライアントは、Databricksのエージェントエンドポイントに対するいたフェースとなります。これは、エンドポイントの呼び出し、レスポンスの解析に関するすべての複雑性に対応するので、UIが行う必要はありません。
# databricks_chat_app/agent_endpoint_client.py
from mlflow.deployments import get_deploy_client
class AgentEndpointClient:
def __init__(self, agent_endpoint_name: str, workspace_client):
"""
Initialize with agent endpoint name and workspace client
The WorkspaceClient handles authentication automatically:
- In Databricks Apps: Uses user's identity
- Local dev: Falls back to DATABRICKS_TOKEN from .env
"""
self.agent_endpoint_name = agent_endpoint_name
self.deploy_client = get_deploy_client("databricks")
mlflow.deployments.get_deploy_client()は、エージェントエンドポイントに対する認証を行うために自動的にWorkspaceClientを使用します。
これでは、predictメソッドを呼び出すことで質問をエージェントに引き渡すことができるようになります。ユーザーがユーザーインタフェースで質問を入力するたびにchatメソッドが呼び出されます。
def chat(self, user_message: str) -> Dict[str, Any]:
"""Send message to agent endpoint and parse response"""
try:
# Call agent via MLflow Deployments Client
# Note: Agent endpoints expect "input" key, not "messages"
response = self.deploy_client.predict(
endpoint=self.agent_endpoint_name,
inputs={
"input": [{"role": "user", "content": user_message}]
}
)
# Parse response to extract summary, charts, and tables
return self._parse_agent_response(response)
エージェントのレスポンスの解析
エージェントからレスポンスを受け取ったら、作成された最終的なエージェントのサマリーと可視化を選択する必要があります。ここでは、以下の情報を解析します:
- エージェントのサマリーテキスト
- チャート(PlotlyのJSON)
- テーブルデータ
エージェントは解析を行うJSONレスポンスを返却するので、UIは結果を表示することができます。我々のリポジトリで解析ロジックを実装しています。ロジックはdatabricks_chat_app/agent_endpoint_client.pyに格納されています。以下は、解析ロジックを駆動させるメインの関数です。
以下を返却します:
- レスポンス - エージェントによって記述された結果のサマリー
- チャート - UIでのレンダリングで使用されるPlotlyチャートのJSON表現
- テーブルデータ - エージェントによって返却される結果データを表現する表形式データ
def _parse_agent_response(self, raw_response: dict) -> Dict[str, Any]:
"""
Parse agent endpoint response to extract summary, charts, and table data
Agent response format: {"object": "response", "output": [...], "id": "..."}
"""
output_array = raw_response.get("output", [])
print(f"[AgentClient] Parsing {len(output_array)} output items...")
return {
"response": self._extract_summary(output_array),
"messages": output_array,
"charts": self._extract_charts(output_array),
"table_data": self._extract_genie_table(output_array),
"error": None
}
エージェントのレスポンスの表示と可視化のレンダリング
これで、エージェントのクライアントを入手し、Streamlit UIがシンプルになりました。これは、単にデータのレンダリングです。初めに、app.pyでエージェントのクライアントを初期化する必要があります。これは、ユーザーを自動で認証するためにWorkspaceClientを活用します。
# databricks_chat_app/app.py
import streamlit as st
import pandas as pd
import plotly.io as pio
import json
from databricks.sdk import WorkspaceClient
from databricks_chat_app.agent_endpoint_client import AgentEndpointClient
import os
# Configure Streamlit
st.set_page_config(
page_title="Data Chatbot",
page_icon="🤖",
layout="wide"
)
# Initialize agent endpoint client
AGENT_ENDPOINT_NAME = os.getenv("AGENT_ENDPOINT_NAME", "my-agent-endpoint")
try:
workspace_client = WorkspaceClient()
agent = AgentEndpointClient(
agent_endpoint_name=AGENT_ENDPOINT_NAME,
workspace_client=workspace_client
)
client_initialized = True
print(f"Connected to agent endpoint: {AGENT_ENDPOINT_NAME}")
except Exception as e:
print(f"Error initializing agent client: {e}")
client_initialized = False
agent = None
エージェントクライアントが初期化されたら、chatメソッドを呼び出して結果を表示する必要があります。エージェントのサマリー、チャート、テーブルが以下で表示されます。
# databricks_chat_app/app.py
# Chat input
if prompt := st.chat_input("What is your question?"):
# Add user message to history
st.session_state.messages.append({
"role": "user",
"content": prompt,
"table_data": None,
"charts": None
})
# Display user message
with st.chat_message("user"):
st.markdown(prompt)
# Call agent and display response
with st.chat_message("assistant"):
thinking_placeholder = st.empty()
thinking_placeholder.markdown("🤖 Agent endpoint is processing...")
# Call agent endpoint
agent_response = agent.chat(prompt)
# Clear thinking indicator
thinking_placeholder.empty()
# Handle errors
if agent_response.get("error"):
response_text = f"**Error:** {agent_response['error']}"
table_data = None
charts = None
else:
# Extract response components
response_text = agent_response.get("response", "Agent processed your request")
table_data = agent_response.get("table_data")
charts = agent_response.get("charts") or []
# Display text response
st.markdown(response_text)
# Display table if present
if table_data:
with st.expander("📊 Table Data", expanded=True):
df = pd.DataFrame(table_data["data"], columns=table_data["columns"])
st.dataframe(df, use_container_width=True)
# Display charts if present
if charts:
for idx, chart_info in enumerate(charts):
with st.expander(f"📈 {chart_info['chart_type'].capitalize()} Chart", expanded=True):
fig = pio.from_json(json.dumps(chart_info["plotly_json"]), skip_invalid=True)
# Style chart
fig.update_layout(
plot_bgcolor='white',
paper_bgcolor='white',
font=dict(color='#000000', size=14)
)
st.plotly_chart(fig, use_container_width=True, key=f"new_chart_{idx}")
# Add assistant message to history
st.session_state.messages.append({
"role": "assistant",
"content": response_text,
"table_data": table_data,
"charts": charts
})
# Rerun to update display
st.rerun()
これによって、以下のスクリーンショットに表示されているように、UIにチャットの履歴とエージェントの結果がレンダリングされます。
結言
インテリジェントなエージェントの構築は、データサイエンティストやAIエキスパートしか手が出せない複雑なものでは無くなっています。Databricksが提供するツールを組み合わせることで、クイックにエージェントを作成し、改善することができます。
関連文書:
Databricks MCPドキュメント: https://docs.databricks.com/aws/ja/generative-ai/mcp/
Unity Catalog AI: https://docs.unitycatalog.io/ai/quickstart/
LangGraphガイド: http://langchain-ai.github.io/langgraph
LangGraph MCPエージェントサンプル: Databricksノートブック
Databricks Genie: https://docs.databricks.com/aws/ja/genie/
Github: https://github.com/bcheng004/plotting-uc-function-managed-mcp