Databricksでブラウザ操作するAIエージェントを作ってデプロイする

# 導入

AIエージェントで多いユースケースにPC操作やブラウザ操作があります。

最近は[MCP](https://github.com/modelcontextprotocol)が盛り上がっており、[Playwright-MCP](https://github.com/microsoft/playwright-mcp)のようなMCPに対応したサーバを利用してエージェントを作成・実行することも多いと思います。

とはいえ、従来のFunction/Tool Callingからブラウザ操作するエージェントもこれはこれで使い勝手が良いと考えています。

というわけで時機を外した感はありますが、Databricksのカスタムエージェントとしてブラウザ操作するエージェントを軽く作ってみたい欲求にかられましたので作成してみます。
また、Mosaic AI Model Serving上にエージェントをデプロイします。

開発はDatabricks on AWS上で行いました。
ノートブックのクラスタはサバ―レスではなく、アクセスモード＝専用のクラスタを利用してください。(Playwrightの依存パッケージがインストールできないため）

# Step1. エージェントを定義する

必要なパッケージをインストールします。
今回、ブラウザ操作は以下の```browser-use```で行いました。
```browser-use```をラップしてエージェントを実装した形となります。

https://github.com/browser-use/browser-use

```python
%pip install -U -qqqq databricks-langchain databricks-agents>=0.16.0 mlflow-skinny[databricks] langgraph==0.3.21 uv loguru rich browser-use nest-asyncio

%restart_python
```

MLflow ChatAgentインターフェースを実装したクラスを定義します。
browser-useのエージェントを内部で実行し、結果を```ChatAgentResponse```クラスに変換して返しています。
LLMはClaude 3.7 Sonnetを利用しました。

```python
%%writefile browser_use_agent.py

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

import mlflow
from databricks_langchain import (
    ChatDatabricks,
)
from langgraph.func import entrypoint, task

from databricks_langchain import ChatDatabricks
from mlflow.pyfunc import ChatAgent
from mlflow.types.agent import (
    ChatAgentMessage,
    ChatAgentResponse,
    ChatContext,
    ChatAgentChunk,
)
import subprocess
import asyncio
from browser_use import Agent, Browser, BrowserConfig
from playwright._impl._driver import compute_driver_executable, get_driver_env


# mlflow tracing
mlflow.langchain.autolog()


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

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

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


@task
async def call_browser_use_agent(llm, task: str):
    """ブラウザを使用してタスクを実行します。"""
    config = BrowserConfig(headless=True, disable_security=True)
    browser = Browser(config=config)

    agent = Agent(
        task=task,
        llm=llm,
        browser=browser,
    )
    results = await agent.run()
    await browser.close()

    return results


@entrypoint()
async def workflow(inputs: dict) -> dict:
    """ブラウザ操作を行うLangGraphのグラフを構築"""

    message = inputs.get("messages", [{}])[-1]
    llm = inputs.get("llm")

    return call_browser_use_agent(llm, message.get("content", None))


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

    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オブジェクト。
        """

        request = {
            "messages": self._convert_messages_to_dict(messages),
            "llm": self.llm,
        }

        results = asyncio.run(workflow.ainvoke(request)).result()

        messages = [
            ChatAgentMessage(
                id=str(uuid.uuid4()),
                role="assistant",
                content=h.result[-1].extracted_content,
                attachments={
                    "url": h.state.url,
                    "title": h.state.title,
                    "screenshot": h.state.screenshot,
                },
            )
            for h in results.history
        ]
        usage = {
            "prompt_tokens": results.total_input_tokens(),
            "completion_tokens": None,
            "total_tokens": results.total_input_tokens(),
        }
        return ChatAgentResponse(
            messages=messages,
            usage=usage,
        )

    def predict_stream(
        self,
        messages: list[ChatAgentMessage],
        context: Optional[ChatContext] = None,
        custom_inputs: Optional[dict[str, Any]] = None,
    ) -> Generator[ChatAgentChunk, None, None]:
        """
        指定されたチャットメッセージリストを使用して、非同期的にエージェントを呼び出し、結果を取得します。

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

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

        request = {
            "messages": self._convert_messages_to_dict(messages),
            "llm": self.llm,
        }
        results = asyncio.run(workflow.ainvoke(request)).result()

        for h in results.history:
            delta = ChatAgentMessage(
                id=str(uuid.uuid4()),
                role="assistant",
                content=h.result[-1].extracted_content,
                attachments={
                    "url": h.state.url,
                    "title": h.state.title,
                    "screenshot": h.state.screenshot,
                },
            )
            yield ChatAgentChunk(
                delta=delta,
                usage={
                    "prompt_tokens": h.metadata.input_tokens,
                    "completion_tokens": None,
                    "total_tokens": h.metadata.input_tokens,
                },
            )


# 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 = BrowserUseChatAgent(llm)
mlflow.models.set_model(AGENT)
```

作成したエージェントをmlflowでロギングします。

```python
import mlflow
from browser_use_agent import LLM_ENDPOINT_NAME
from mlflow.models.resources import DatabricksServingEndpoint
from rich import print

resources = [DatabricksServingEndpoint(endpoint_name=LLM_ENDPOINT_NAME)]

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

with mlflow.start_run():
    logged_agent_info = mlflow.pyfunc.log_model(
        artifact_path="agent",
        python_model="browser_use_agent.py",
        input_example=input_example,
        pip_requirements=[
            "mlflow",
            "langgraph==0.3.21",
            "databricks-langchain==0.4.1",
            "browser-use==0.1.40",

        ],
        resources=resources,
    )
```

ロギングしたモデルを利用して推論してみます。

```python
run_id = logged_agent_info.run_id
mlflow.models.predict(
    model_uri=f"runs:/{run_id}/agent",
    input_data={"messages": [{"role": "user", "content": "Hello!"}]},
    env_manager="uv",
)
```

```:出力結果
INFO     [agent] 🚀 Starting task: Hello!
INFO     [agent] 📍 Step 1
INFO     [agent] 🤷 Eval: Unknown - This is the first step of the task. I am starting with a blank page.
INFO     [agent] 🧠 Memory: I am starting with a blank page. The ultimate task is to say 'Hello!'. This is a simple greeting task that I can complete immediately.
INFO     [agent] 🎯 Next goal: Complete the task by saying Hello!
INFO     [agent] 🛠️  Action 1/1: {"done":{"text":"Hello!","success":true}}
INFO     [agent] 📄 Result: Hello!
INFO     [agent] ✅ Task completed
INFO     [agent] ✅ Successfully
{"messages": [{"role": "assistant", "content": "Hello!", "id": "89fc1657-3744-42c2-8b36-2b66fdb59921", "attachments": {"url": "about:blank", "title": "", "screenshot":(省略)
```

ブラウザ操作はおこなっていませんが、問題なく動作はしていそうです。



# Step2. Mosaic AI Agent Frameworkを使ってデプロイ

では、Mosaic AI Model Serving上にエージェントをデプロイしてみましょう。

まず、Step1でロギングしたエージェントをUnity Catalog上に登録します。

```python
import mlflow
mlflow.set_registry_uri("databricks-uc")

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

# Unity Catalogへの登録
uc_registered_model_info = mlflow.register_model(
    model_uri=f"runs:/{run_id}/agent", name=UC_MODEL_NAME
)
```

次にMosaic AI Agent Frameworkを使ってMosaic AI Model Servingへデプロイします。

```python
from databricks import agents

agents.deploy(
    model_name=UC_MODEL_NAME,
    model_version=uc_registered_model_info.version,
    scale_to_zero=True,
)
```

問題なければ数分程度でデプロイが完了します。

:::note warn
たまにPlaywrightの依存関係インストールに失敗することがありました。
うまくいかない場合はデプロイをやり直してください。
:::


![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/c69c16ef-52b5-434b-844c-2ccf5e36a5c9.png)

# Step3. Playgroundで使ってみる

では、デプロイしたエージェントを試しに使ってみます。

今回はDatabricksのPlayground上でエージェントに指示を出してみました。

まず、「日経平均株価を教えて」もらいます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/f184de6c-f1f9-4fd0-9f3f-7c4bc6e99a0b.png)

操作している画面は表示されませんが、Googleで検索して情報を取得してくれました。

検索系ばかりですが、もう１個ぐらい指示してみます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/c4974f15-bb23-45fe-9f6f-d3cd79cd0f77.png)

Bリーグ公式サイトの中まで入って情報を取って来てくれていますね。

なお、デプロイしたエージェントは、ブラウザ操作画面のスクリーンショットも出力するようにしていますので、それを利用してブラウザ操作画面履歴を出したりといったこともできるようにしています。（通信量が増えてしまっていますが）

# まとめ

習作として、ブラウザ操作エージェントをDatabricks上で作成&デプロイしてみました。
作りが甘く実用性はまだまだです。特に結果を全てassistantとして返したり、ストリーミング処理が適当だったりは改善の余地が大きい。

とはいえ、ちょっとしたブラウザ処理の自動化には利用できるかなと思います。
（個人的にはマルチエージェントシステム内のひとつのエージェントとして使えればと考えています）

次はMCPに対応したエージェントも作ってみたいなあ。