DatabricksでPlaywright-MCPと連携するChatAgentを作ってデプロイする #LLM – Qiita

%pip install langchain-mcp-adapters langgraph databricks-langchain databricks-agents uv rich nest-asyncio playwright 'nodejs-bin[cmd]'
%pip install --upgrade 'mlflow>=3.0.0rc0' --pre

%restart_python

import nest_asyncio

nest_asyncio.apply()

%%writefile playwright_mcp_agent.py

from typing import Literal, Generator, List, Optional, Any, Dict, Mapping, Union
import uuid
import subprocess
import asyncio
import sys

import mlflow
from databricks_langchain import ChatDatabricks
from langgraph.prebuilt import create_react_agent

from langchain_core.messages import BaseMessage, convert_to_openai_messages
from mlflow.pyfunc import ChatAgent
from mlflow.types.agent import (
    ChatAgentChunk,
    ChatAgentMessage,
    ChatAgentResponse,
    ChatContext,
)
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from langchain_mcp_adapters.tools import load_mcp_tools
from playwright._impl._driver import compute_driver_executable, get_driver_env
from functools import reduce

# Tracingの有効化
mlflow.langchain.autolog(log_models=False)

def install_playwright():
    """Playwrightの依存関係をインストールします。"""

    driver_executable, driver_cli = compute_driver_executable()
    args = [driver_executable, driver_cli, "install", "firefox", "--with-deps"]

    proc = subprocess.run(
        args, env=get_driver_env(), capture_output=True, text=True, check=True
    )
    return proc == 0

async def invoke_with_mcp_tools(model, server_params, messages: list):
    """MCP Serverと通信して回答を取得する"""

    async with stdio_client(server_params, errlog=None) as (read, write):
        async with ClientSession(read, write) as session:
            # 接続を初期化
            await session.initialize()

            # ツールを取得
            tools = await load_mcp_tools(session)

            # エージェントを作成して実行
            agent = create_react_agent(model, tools)
            return await agent.ainvoke({"messages": messages})

class MCPChatAgent(ChatAgent):
    def __init__(self, model):
        """LangGraphのグラフを指定して初期化"""
        pass
        self.model = model

    def predict(
        self,
        messages: list[ChatAgentMessage],
        context: Optional[ChatContext] = None,
        custom_inputs: Optional[dict[str, Any]] = None,
    ) -> ChatAgentResponse:
        """
        指定されたチャットメッセージリストを使用して回答を生成する

        Args:
            messages (list[ChatAgentMessage]): チャットエージェントメッセージのリスト。
            context (Optional[ChatContext]): オプションのチャットコンテキスト。
            custom_inputs (Optional[dict[str, Any]]): カスタム入力のオプション辞書。

        Returns:
            ChatAgentResponse: 予測結果を含むChatAgentResponseオブジェクト。
        """

        server_params = StdioServerParameters(
            command="npx",
            args=["@playwright/mcp@latest", "--headless", "--browser", "firefox"],
        )

        result = asyncio.run(
            invoke_with_mcp_tools(
                self.model,
                server_params,
                self._convert_messages_to_dict(messages),
            )
        )

        messages = []
        usages = []
        for msg in result.get("messages"):
            messages.append(self._convert_lc_message_to_chat_message(msg))

            if msg.response_metadata:
                usages.append(msg.response_metadata.get("usage"))

        usage = self._sum_usages(usages)
        return ChatAgentResponse(messages=messages, usage=usage)

    def _convert_lc_message_to_chat_message(
        self, lc_message: BaseMessage
    ) -> ChatAgentMessage:
        """LangChainメッセージをChatAgentMessageに変換する。"""
        msg = convert_to_openai_messages(lc_message)
        if not "id" in msg:
            msg.update({"id": str(uuid.uuid4())})

        return ChatAgentMessage(**msg)

    def _sum_usages(self, usages: list[dict]) -> dict:
        """使用量のリストから使用量を合計する。"""

        def add_usages(a: dict, b: dict) -> dict:
            pt = "prompt_tokens"
            ct = "completion_tokens"
            tt = "total_tokens"
            return {
                pt: a.get(pt, 0) + b.get(pt, 0),
                ct: a.get(ct, 0) + b.get(ct, 0),
                tt: a.get(tt, 0) + b.get(tt, 0),
            }

        return reduce(add_usages, usages, {})

# PlaywrightをInstall
install_playwright()

# DatabricksネイティブのClaude 3.7 SonnetをLLMとして利用
LLM_ENDPOINT_NAME = "databricks-claude-3-7-sonnet"
llm = ChatDatabricks(model=LLM_ENDPOINT_NAME)

# エージェント作成
AGENT = MCPChatAgent(model=llm)
mlflow.models.set_model(AGENT)

import mlflow
from playwright_mcp_agent import LLM_ENDPOINT_NAME
from mlflow.models.resources import DatabricksServingEndpoint

resources = [DatabricksServingEndpoint(endpoint_name=LLM_ENDPOINT_NAME)]

input_example = {
    "messages": [
        {
            "role": "user",
            "content": "今日の日経平均株価はいくら？"
        }
    ]
}

with mlflow.start_run():
    logged_agent_info = mlflow.pyfunc.log_model(
        # artifact_path ="playwright_mcp_agent", # for mlflow 2.0
        name="playwright_mcp_agent", # for mlflow 3.0
        python_model="playwright_mcp_agent.py",
        input_example=input_example,
        pip_requirements=[
            "mlflow",
            "langgraph==0.3.29",
            "langchain-mcp-adapters==0.0.7",
            "databricks-langchain==0.4.2",
            "playwright==1.51.0",
            "nodejs-bin[cmd]",
        ],
        resources=resources,
        model_type="agent", # for mlflow 3.0
    )

catalog = "training"
schema = "llm"
model_name = "playwright_mcp_agent"
UC_MODEL_NAME = f"{catalog}.{schema}.{model_name}"

mlflow.set_registry_uri("databricks-uc")

# モデルをUCに登録する
uc_registered_model_info = mlflow.register_model(
    model_uri=logged_model.model_uri, name=UC_MODEL_NAME
)

from databricks import agents
agents.deploy(UC_MODEL_NAME, uc_registered_model_info.version, scale_to_zero=True)