Merge remote-tracking branch 'upstream/main' into stream_function

Author: katiue

src/agent/custom_agent.py

@@ -9,6 +9,10 @@
 import pdb
 import traceback
 from typing import Optional, Type
+from PIL import Image, ImageDraw, ImageFont
+import os
+import base64
+import io
 
 from browser_use.agent.prompts import SystemPrompt
 from browser_use.agent.service import Agent
@@ -227,6 +231,119 @@ async def step(self, step_info: Optional[CustomAgentStepInfo] = None) -> None:
                     )
             if state:
                 self._make_history_item(model_output, state, result)
+    def create_history_gif(
+            self,
+            output_path: str = 'agent_history.gif',
+            duration: int = 3000,
+            show_goals: bool = True,
+            show_task: bool = True,
+            show_logo: bool = False,
+            font_size: int = 40,
+            title_font_size: int = 56,
+            goal_font_size: int = 44,
+            margin: int = 40,
+            line_spacing: float = 1.5,
+    ) -> None:
+        """Create a GIF from the agent's history with overlaid task and goal text."""
+        if not self.history.history:
+            logger.warning('No history to create GIF from')
+            return
+
+        images = []
+        # if history is empty or first screenshot is None, we can't create a gif
+        if not self.history.history or not self.history.history[0].state.screenshot:
+            logger.warning('No history or first screenshot to create GIF from')
+            return
+
+        # Try to load nicer fonts
+        try:
+            # Try different font options in order of preference
+            font_options = ['Helvetica', 'Arial', 'DejaVuSans', 'Verdana']
+            font_loaded = False
+
+            for font_name in font_options:
+                try:
+                    import platform
+                    if platform.system() == "Windows":
+                        # Need to specify the abs font path on Windows
+                        font_name = os.path.join(os.getenv("WIN_FONT_DIR", "C:\\Windows\\Fonts"), font_name + ".ttf")
+                    regular_font = ImageFont.truetype(font_name, font_size)
+                    title_font = ImageFont.truetype(font_name, title_font_size)
+                    goal_font = ImageFont.truetype(font_name, goal_font_size)
+                    font_loaded = True
+                    break
+                except OSError:
+                    continue
+
+            if not font_loaded:
+                raise OSError('No preferred fonts found')
+
+        except OSError:
+            regular_font = ImageFont.load_default()
+            title_font = ImageFont.load_default()
+
+            goal_font = regular_font
+
+        # Load logo if requested
+        logo = None
+        if show_logo:
+            try:
+                logo = Image.open('./static/browser-use.png')
+                # Resize logo to be small (e.g., 40px height)
+                logo_height = 150
+                aspect_ratio = logo.width / logo.height
+                logo_width = int(logo_height * aspect_ratio)
+                logo = logo.resize((logo_width, logo_height), Image.Resampling.LANCZOS)
+            except Exception as e:
+                logger.warning(f'Could not load logo: {e}')
+
+        # Create task frame if requested
+        if show_task and self.task:
+            task_frame = self._create_task_frame(
+                self.task,
+                self.history.history[0].state.screenshot,
+                title_font,
+                regular_font,
+                logo,
+                line_spacing,
+            )
+            images.append(task_frame)
+
+        # Process each history item
+        for i, item in enumerate(self.history.history, 1):
+            if not item.state.screenshot:
+                continue
+
+            # Convert base64 screenshot to PIL Image
+            img_data = base64.b64decode(item.state.screenshot)
+            image = Image.open(io.BytesIO(img_data))
+
+            if show_goals and item.model_output:
+                image = self._add_overlay_to_image(
+                    image=image,
+                    step_number=i,
+                    goal_text=item.model_output.current_state.thought,
+                    regular_font=regular_font,
+                    title_font=title_font,
+                    margin=margin,
+                    logo=logo,
+                )
+
+            images.append(image)
+
+        if images:
+            # Save the GIF
+            images[0].save(
+                output_path,
+                save_all=True,
+                append_images=images[1:],
+                duration=duration,
+                loop=0,
+                optimize=False,
+            )
+            logger.info(f'Created GIF at {output_path}')
+        else:
+            logger.warning('No images found in history to create GIF')
 
     async def run(self, max_steps: int = 100) -> AgentHistoryList:
         """Execute the task with maximum number of steps"""
@@ -283,3 +400,6 @@ async def run(self, max_steps: int = 100) -> AgentHistoryList:
 
             if not self.injected_browser and self.browser:
                 await self.browser.close()
+
+            if self.generate_gif:
+                self.create_history_gif()

src/browser/custom_browser.py

@@ -123,4 +123,4 @@ async def _setup_browser(self, playwright: Playwright) -> PlaywrightBrowser:
                 return browser
             except Exception as e:
                 logger.error(f'Failed to initialize Playwright browser: {str(e)}')
-                raise
+                raise

tests/test_browser_use.py

@@ -3,6 +3,7 @@
 # @Author  : wenshao
 # @ProjectName: browser-use-webui
 # @FileName: test_browser_use.py
+import pdb
 
 from dotenv import load_dotenv
 
@@ -28,20 +29,29 @@ async def test_browser_use_org():
         BrowserContextWindowSize,
     )
 
+    # llm = utils.get_llm_model(
+    #     provider="azure_openai",
+    #     model_name="gpt-4o",
+    #     temperature=0.8,
+    #     base_url=os.getenv("AZURE_OPENAI_ENDPOINT", ""),
+    #     api_key=os.getenv("AZURE_OPENAI_API_KEY", ""),
+    # )
+
     llm = utils.get_llm_model(
-        provider="azure_openai",
-        model_name="gpt-4o",
-        temperature=0.8,
-        base_url=os.getenv("AZURE_OPENAI_ENDPOINT", ""),
-        api_key=os.getenv("AZURE_OPENAI_API_KEY", ""),
+        provider="deepseek",
+        model_name="deepseek-chat",
+        temperature=0.8
     )
 
     window_w, window_h = 1920, 1080
+    use_vision = False
+    chrome_path = os.getenv("CHROME_PATH", None)
 
     browser = Browser(
         config=BrowserConfig(
             headless=False,
             disable_security=True,
+            chrome_instance_path=chrome_path,
             extra_chromium_args=[f"--window-size={window_w},{window_h}"],
         )
     )
@@ -59,6 +69,7 @@ async def test_browser_use_org():
             task="go to google.com and type 'OpenAI' click search and give me the first url",
             llm=llm,
             browser_context=browser_context,
+            use_vision=use_vision
         )
         history: AgentHistoryList = await agent.run(max_steps=10)
 
@@ -208,6 +219,122 @@ async def test_browser_use_custom():
         await browser.close()
 
 
+async def test_browser_use_custom_v2():
+    from browser_use.browser.context import BrowserContextWindowSize
+    from browser_use.browser.browser import BrowserConfig
+    from playwright.async_api import async_playwright
+
+    from src.agent.custom_agent import CustomAgent
+    from src.agent.custom_prompts import CustomSystemPrompt
+    from src.browser.custom_browser import CustomBrowser
+    from src.browser.custom_context import BrowserContextConfig
+    from src.controller.custom_controller import CustomController
+
+    window_w, window_h = 1920, 1080
+
+    # llm = utils.get_llm_model(
+    #     provider="azure_openai",
+    #     model_name="gpt-4o",
+    #     temperature=0.8,
+    #     base_url=os.getenv("AZURE_OPENAI_ENDPOINT", ""),
+    #     api_key=os.getenv("AZURE_OPENAI_API_KEY", ""),
+    # )
+
+    # llm = utils.get_llm_model(
+    #     provider="gemini",
+    #     model_name="gemini-2.0-flash-exp",
+    #     temperature=1.0,
+    #     api_key=os.getenv("GOOGLE_API_KEY", "")
+    # )
+
+    llm = utils.get_llm_model(
+        provider="deepseek",
+        model_name="deepseek-chat",
+        temperature=0.8
+    )
+
+    # llm = utils.get_llm_model(
+    #     provider="ollama", model_name="qwen2.5:7b", temperature=0.8
+    # )
+
+    controller = CustomController()
+    use_own_browser = True
+    disable_security = True
+    use_vision = False  # Set to False when using DeepSeek
+    tool_call_in_content = True  # Set to True when using Ollama
+    max_actions_per_step = 1
+    playwright = None
+    browser = None
+    browser_context = None
+
+    try:
+        if use_own_browser:
+            chrome_path = os.getenv("CHROME_PATH", None)
+            if chrome_path == "":
+                chrome_path = None
+        else:
+            chrome_path = None
+        browser = CustomBrowser(
+            config=BrowserConfig(
+                headless=False,
+                disable_security=disable_security,
+                chrome_instance_path=chrome_path,
+                extra_chromium_args=[f"--window-size={window_w},{window_h}"],
+            )
+        )
+        browser_context = await browser.new_context(
+            config=BrowserContextConfig(
+                trace_path="./tmp/traces",
+                save_recording_path="./tmp/record_videos",
+                no_viewport=False,
+                browser_window_size=BrowserContextWindowSize(
+                    width=window_w, height=window_h
+                ),
+            )
+        )
+        agent = CustomAgent(
+            task="go to google.com and type 'OpenAI' click search and give me the first url",
+            add_infos="",  # some hints for llm to complete the task
+            llm=llm,
+            browser=browser,
+            browser_context=browser_context,
+            controller=controller,
+            system_prompt_class=CustomSystemPrompt,
+            use_vision=use_vision,
+            tool_call_in_content=tool_call_in_content,
+            max_actions_per_step=max_actions_per_step
+        )
+        history: AgentHistoryList = await agent.run(max_steps=10)
+
+        print("Final Result:")
+        pprint(history.final_result(), indent=4)
+
+        print("\nErrors:")
+        pprint(history.errors(), indent=4)
+
+        # e.g. xPaths the model clicked on
+        print("\nModel Outputs:")
+        pprint(history.model_actions(), indent=4)
+
+        print("\nThoughts:")
+        pprint(history.model_thoughts(), indent=4)
+        # close browser
+    except Exception:
+        import traceback
+
+        traceback.print_exc()
+    finally:
+        # 显式关闭持久化上下文
+        if browser_context:
+            await browser_context.close()
+
+        # 关闭 Playwright 对象
+        if playwright:
+            await playwright.stop()
+        if browser:
+            await browser.close()
+
 if __name__ == "__main__":
     # asyncio.run(test_browser_use_org())
-    asyncio.run(test_browser_use_custom())
+    # asyncio.run(test_browser_use_custom())
+    asyncio.run(test_browser_use_custom_v2())