apunkt/MetaGPT
1
0
Fork 0
mirror of https://github.com/FoundationAgents/MetaGPT.git synced 2026-05-24 14:15:17 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix code issues

Browse source
This commit is contained in:
黄伟韬 2024-09-12 13:26:35 +08:00
parent d9dacb6305
commit 65dc47c90d
2 changed files with 30 additions and 33 deletions
Download patch file Download diff file Expand all files Collapse all files

61
metagpt/tools/libs/image_getter.py
View file

@ -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."