[browser_use] Get prompt and response through callback

browser_use_callback.py

"""
第二个demo.
"""
from typing import Any, Optional
from uuid import UUID
from langchain_core.callbacks import StdOutCallbackHandler
from langchain_core.callbacks.base import BaseCallbackHandler
from langchain_core.outputs import LLMResult
from langchain_openai import ChatOpenAI
from browser_use import Agent, BrowserConfig, Browser
from browser_use.browser.views import BrowserState
from browser_use.agent.views import AgentOutput

import time
import os
import contextvars
import logging
import asyncio
import json
from dotenv import load_dotenv
load_dotenv()

RUN_ID = time.strftime("%Y%m%d-%H%M%S")
LOG_DIR = f'log/{RUN_ID}'

logger = logging.getLogger(__name__)
step_var = contextvars.ContextVar('step', default=0)


def json_dump(obj, fn):
    with open(fn, 'w') as fp:
        json.dump(obj, fp, indent=1, ensure_ascii=False, separators=(',', ':'))

class E0CallbackHandler(BaseCallbackHandler):
    def on_llm_start(self, serialized: dict[str, Any], prompts: list[str], *, run_id: UUID, parent_run_id: UUID | None = None, tags: list[str] | None = None, metadata: dict[str, Any] | None = None, **kwargs: Any) -> Any:
        step = step_var.get()
        logger.info(f'ON_LLM_START run_id:{run_id} step:{step}')
        json_dump(prompts, f"{LOG_DIR}/{step:04d}-a.json")

    def on_llm_end(self, response: LLMResult, *,
        run_id: UUID, parent_run_id: Optional[UUID] = None, **kwargs: Any):
        step = step_var.get()
        logger.info(f'ON_LLM_END run_id:{run_id} step:{step}')
        json_dump(response.model_dump(),
            f'{LOG_DIR}/{step:04d}-b.json')

async def agent_step_callback(state:BrowserState, output:AgentOutput, step:int) -> None:
    step_var.set(step)
    logger.info(f'__CALLBACK__ __STEP__ {step}')

async def main():
    os.makedirs(LOG_DIR, exist_ok=True)
    config = BrowserConfig(
        cdp_url="http://localhost:9333",
        headless=True,
    )
    browser = Browser(config)

    handler = E0CallbackHandler()
    agent = Agent(
        browser=browser,
        use_vision=False,
        task="Compare the price of gpt-4o and DeepSeek-V3",
        register_new_step_callback=agent_step_callback,
        llm=ChatOpenAI(
            model="gpt-4o-0513-global",
            callbacks=[handler])
    )
    await agent.run()
    print(f'RUN_ID: {RUN_ID}')
    print(f'LOG_DIR: {LOG_DIR}')

asyncio.run(main())