インモビインモビ

%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

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

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

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":(省略)

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
)

from databricks import agents

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