From 65dc47c90d105b5ef93b59ff9db94ca92492211e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=BB=84=E4=BC=9F=E9=9F=AC?= <huangweitao@fuzhi.ai>
Date: Thu, 12 Sep 2024 13:26:35 +0800
Subject: [PATCH] fix code issues

---
 metagpt/strategy/experience_retriever.py |  2 +-
 metagpt/tools/libs/image_getter.py       | 61 +++++++++++-------------
 2 files changed, 30 insertions(+), 33 deletions(-)

diff --git a/metagpt/strategy/experience_retriever.py b/metagpt/strategy/experience_retriever.py
index 5cb0573b2..e974e41a0 100644
--- a/metagpt/strategy/experience_retriever.py
+++ b/metagpt/strategy/experience_retriever.py
@@ -845,7 +845,7 @@ Consider this example only after you have obtained the content of system design
 Suppose the system design and project schedule prescribes three files index.html, style.css, script.js, to follow the design and schedule, I will create a plan consisting of three tasks, each corresponding to the creation of one of the required files: `index.html`, `style.css`, and `script.js`. 
 
 Here's the plan:
-
+[Optional] 0. **Task 0**: Obtain images before coding.
 1. **Task 1**: Create `index.html` - This file will contain the HTML structure necessary for the game's UI.
 2. **Task 2**: Create `style.css` - This file will define the CSS styles to make the game visually appealing and responsive.
 3. **Task 3**: Create `script.js` - This file will contain the JavaScript code for the game logic and UI interactions.
diff --git a/metagpt/tools/libs/image_getter.py b/metagpt/tools/libs/image_getter.py
index 1420f9d7f..a8f982fa5 100644
--- a/metagpt/tools/libs/image_getter.py
+++ b/metagpt/tools/libs/image_getter.py
@@ -1,9 +1,5 @@
 from __future__ import annotations
 
-import base64
-import os
-import re
-from pathlib import Path
 from typing import Optional
 
 from playwright.async_api import Browser as Browser_
@@ -11,9 +7,28 @@ from playwright.async_api import BrowserContext, Page, Playwright, async_playwri
 from pydantic import BaseModel, ConfigDict, Field
 
 from metagpt.tools.tool_registry import register_tool
+from metagpt.utils.common import decode_image
 from metagpt.utils.proxy_env import get_proxy_from_env
 from metagpt.utils.report import BrowserReporter
 
+DOWNLOAD_PICTURE_JAVASCRIPT = """
+async () => {{
+    var img = document.querySelector('{img_element_selector}');
+    if (img && img.src) {{
+        const response = await fetch(img.src);
+        if (response.ok) {{
+            const blob = await response.blob();
+            return await new Promise(resolve => {{
+                const reader = new FileReader();
+                reader.onloadend = () => resolve(reader.result);
+                reader.readAsDataURL(blob);
+            }});
+        }}
+    }}
+    return null;
+}}
+"""
+
 
 @register_tool(include_functions=["get_image"])
 class ImageGetter(BaseModel):
@@ -30,6 +45,8 @@ class ImageGetter(BaseModel):
     headless: bool = Field(default=True)
     proxy: Optional[dict] = Field(default_factory=get_proxy_from_env)
     reporter: BrowserReporter = Field(default_factory=BrowserReporter)
+    url: str = "https://unsplash.com/s/photos/{search_term}/"
+    img_element_selector: str = ".zNNw1 > div > img:nth-of-type(2)"
 
     async def start(self) -> None:
         """Starts Playwright and launches a browser"""
@@ -44,45 +61,25 @@ class ImageGetter(BaseModel):
         Get an image related to the search term.
 
         Args:
-            search_term (str): The term to search for the image.
+            search_term (str): The term to search for the image. The search term must be in English; using any other language may lead to a mismatch.
             image_save_path (str): The file path where the image will be saved.
         """
         # Search for images from https://unsplash.com/s/photos/
-        url = f"https://unsplash.com/s/photos/{search_term}/"
+
         if self.page is None:
             await self.start()
-        await self.page.goto(url, wait_until="domcontentloaded")
+        await self.page.goto(self.url.format(search_term=search_term), wait_until="domcontentloaded")
         # Wait until the image element is loaded
         try:
-            await self.page.wait_for_selector(".zNNw1 > div > img:nth-of-type(2)")
+            await self.page.wait_for_selector(self.img_element_selector)
         except TimeoutError:
             return f"{search_term} not found. Please broaden the search term."
         # Get the base64 code of the first  retrieved image
         image_base64 = await self.page.evaluate(
-            """async () => {
-            var img = document.querySelector('.zNNw1 > div > img:nth-of-type(2)');
-            if (img && img.src) {
-                const response = await fetch(img.src);
-                if (response.ok) {
-                    const blob = await response.blob();
-                    return await new Promise(resolve => {
-                        const reader = new FileReader();
-                        reader.onloadend = () => resolve(reader.result);
-                        reader.readAsDataURL(blob);
-                    });
-                }
-            }
-            return null;
-        }"""
+            DOWNLOAD_PICTURE_JAVASCRIPT.format(img_element_selector=self.img_element_selector)
         )
         if image_base64:
-            # Save image
-            file_path = Path(image_save_path)
-            os.makedirs(file_path.parent, exist_ok=True)
-            with open(image_save_path, "wb") as f:
-                imgstr = re.sub("data:image/.*?;base64,", "", image_base64)
-                image_data = base64.b64decode(imgstr)
-                f.write(image_data)
+            image = decode_image(image_base64)
+            image.save(image_save_path)
             return f"{search_term} found. The image is saved in {image_save_path}."
-        else:
-            return f"{search_term} not found. Please broaden the search term."
+        return f"{search_term} not found. Please broaden the search term."