From 1938a625a3300e9bd6a96e8414cd3880882d819c Mon Sep 17 00:00:00 2001
From: BespalovSergey <50576059+BespalovSergey@users.noreply.github.com>
Date: Thu, 4 Jul 2024 14:27:33 +0300
Subject: [PATCH] Html rendering tool (#59)

* add HtmlRenderTool

* remove HtmlRenderTool from tools init

* add HtmlRenderTool to tools init

* add HtmlRenderTool tests

* update PythonLinterTool

* minor improvements

---------

Co-authored-by: User <serfintw@yadnex.ru>
Co-authored-by: whimo <whimox@gmail.com>
---
 motleycrew/common/exceptions.py           |  30 +++--
 motleycrew/tools/__init__.py              |   7 +-
 motleycrew/tools/code/python_linter.py    |   7 +-
 motleycrew/tools/html_render_tool.py      | 150 ++++++++++++++++++++++
 pytest.ini                                |   4 +
 tests/test_tools/test_html_render_tool.py |  25 ++++
 6 files changed, 207 insertions(+), 16 deletions(-)
 create mode 100644 motleycrew/tools/html_render_tool.py
 create mode 100644 pytest.ini
 create mode 100644 tests/test_tools/test_html_render_tool.py

diff --git a/motleycrew/common/exceptions.py b/motleycrew/common/exceptions.py
index e1fe31f2..7eaefba1 100644
--- a/motleycrew/common/exceptions.py
+++ b/motleycrew/common/exceptions.py
@@ -1,16 +1,18 @@
 """ Module description"""
+
 from typing import Any, Optional
 
 from motleycrew.common import Defaults
 
 
 class LLMFamilyNotSupported(Exception):
-    """ Description
+    """Description
 
     Args:
         llm_framework (str):
         llm_family (str):
     """
+
     def __init__(self, llm_framework: str, llm_family: str):
         self.llm_framework = llm_framework
         self.llm_family = llm_family
@@ -21,7 +23,7 @@ def __str__(self) -> str:
 
 class LLMFrameworkNotSupported(Exception):
     def __init__(self, llm_framework: str):
-        """ Description
+        """Description
 
         Args:
             llm_framework (str):
@@ -33,11 +35,12 @@ def __str__(self) -> str:
 
 
 class AgentNotMaterialized(Exception):
-    """ Description
+    """Description
+
+    Args:
+        agent_name (str):
+    """
 
-   Args:
-       agent_name (str):
-   """
     def __init__(self, agent_name: str):
         self.agent_name = agent_name
 
@@ -46,11 +49,12 @@ def __str__(self) -> str:
 
 
 class CannotModifyMaterializedAgent(Exception):
-    """ Description
+    """Description
 
     Args:
         agent_name (str):
     """
+
     def __init__(self, agent_name: str | None):
         self.agent_name = agent_name
 
@@ -65,7 +69,7 @@ class TaskDependencyCycleError(Exception):
 
 
 class IntegrationTestException(Exception):
-    """ Integration tests exception
+    """Integration tests exception
 
     Args:
         test_names (list[str]): list of names of failed integration tests
@@ -79,18 +83,21 @@ def __str__(self):
 
 
 class IpynbIntegrationTestResultNotFound(Exception):
-    """ Ipynb integration test not found result file exception
+    """Ipynb integration test not found result file exception
 
     Args:
         ipynb_path (str): path to running ipynb
         result_path (str): path to execution result file
     """
+
     def __init__(self, ipynb_path: str, result_path: str):
         self.ipynb_path = ipynb_path
         self.result_path = result_path
 
     def __str__(self):
-        return "File result {} of the ipynb {} execution, not found.".format(self.result_path, self.ipynb_path)
+        return "File result {} of the ipynb {} execution, not found.".format(
+            self.result_path, self.ipynb_path
+        )
 
 
 class ModuleNotInstalled(Exception):
@@ -111,7 +118,7 @@ def __str__(self):
         msg = "{} is not installed".format(self.module_name)
 
         if self.install_command is not None:
-            msg = "{}, {}".format(msg, self.install_command)
+            msg = "{}, please install ({})".format(msg, self.install_command)
 
         return "{}.".format(msg)
 
@@ -133,4 +140,5 @@ def __str__(self):
 
 class InvalidOutput(Exception):
     """Raised in output handlers when an agent's output is not accepted"""
+
     pass
diff --git a/motleycrew/tools/__init__.py b/motleycrew/tools/__init__.py
index 8500ae6c..52f3f6a0 100644
--- a/motleycrew/tools/__init__.py
+++ b/motleycrew/tools/__init__.py
@@ -1,9 +1,10 @@
-from .tool import MotleyTool
+from motleycrew.tools.tool import MotleyTool
 
 from .autogen_chat_tool import AutoGenChatTool
+from .code.postgresql_linter import PostgreSQLLinterTool
+from .code.python_linter import PythonLinterTool
+from .html_render_tool import HTMLRenderTool
 from .image.dall_e import DallEImageGeneratorTool
 from .llm_tool import LLMTool
 from .mermaid_evaluator_tool import MermaidEvaluatorTool
 from .python_repl import PythonREPLTool
-from .code.postgresql_linter import PostgreSQLLinterTool
-from .code.python_linter import PythonLinterTool
diff --git a/motleycrew/tools/code/python_linter.py b/motleycrew/tools/code/python_linter.py
index 2c764db4..b70c967c 100644
--- a/motleycrew/tools/code/python_linter.py
+++ b/motleycrew/tools/code/python_linter.py
@@ -4,11 +4,14 @@
 from langchain_core.tools import StructuredTool
 from langchain_core.pydantic_v1 import BaseModel, Field
 
+try:
+    from aider.linter import Linter
+except ImportError:
+    Linter = None
+
 from motleycrew.common.utils import ensure_module_is_installed
 from motleycrew.tools import MotleyTool
 
-Linter = None
-
 
 class PythonLinterTool(MotleyTool):
 
diff --git a/motleycrew/tools/html_render_tool.py b/motleycrew/tools/html_render_tool.py
new file mode 100644
index 00000000..f730b318
--- /dev/null
+++ b/motleycrew/tools/html_render_tool.py
@@ -0,0 +1,150 @@
+from datetime import datetime
+from pathlib import Path
+from typing import Tuple, Optional
+
+from motleycrew.common.utils import ensure_module_is_installed
+
+try:
+    from selenium import webdriver
+    from selenium.webdriver.chrome.service import Service
+except ImportError:
+    webdriver = None
+    Service = None
+
+from langchain.tools import Tool
+from langchain_core.pydantic_v1 import BaseModel, Field
+
+from motleycrew.tools import MotleyTool
+from motleycrew.common import logger
+
+
+class HTMLRenderer:
+    def __init__(
+        self,
+        work_dir: str,
+        executable_path: str | None = None,
+        headless: bool = True,
+        window_size: Optional[Tuple[int, int]] = None,
+    ):
+        """Helper for rendering HTML code as an image"""
+        ensure_module_is_installed(
+            "selenium",
+            "see documentation: https://pypi.org/project/selenium/, ChromeDriver is also required",
+        )
+
+        self.work_dir = Path(work_dir).resolve()
+        self.html_dir = self.work_dir / "html"
+        self.images_dir = self.work_dir / "images"
+
+        self.options = webdriver.ChromeOptions()
+        if headless:
+            self.options.add_argument("--headless")
+        self.service = Service(executable_path=executable_path)
+
+        self.window_size = window_size
+
+    def render_image(self, html: str, file_name: str | None = None):
+        """Create image with png extension from html code
+
+        Args:
+            html (str): html code for rendering image
+            file_name (str): file name with not extension
+        Returns:
+            file path to created image
+        """
+        logger.info("Trying to render image from HTML code")
+        html_path, image_path = self.build_save_file_paths(file_name)
+        browser = webdriver.Chrome(options=self.options, service=self.service)
+        try:
+            if self.window_size:
+                logger.info("Setting window size to {}".format(self.window_size))
+                browser.set_window_size(*self.window_size)
+
+            url = "data:text/html;charset=utf-8,{}".format(html)
+            browser.get(url)
+
+            logger.info("Taking screenshot")
+            is_created_img = browser.get_screenshot_as_file(image_path)
+        finally:
+            browser.close()
+            browser.quit()
+
+        if not is_created_img:
+            logger.error("Failed to render image from HTML code {}".format(image_path))
+            return "Failed to render image from HTML code"
+
+        with open(html_path, "w") as f:
+            f.write(html)
+        logger.info("Saved the HTML code to {}".format(html_path))
+        logger.info("Saved the rendered HTML screenshot to {}".format(image_path))
+
+        return image_path
+
+    def build_save_file_paths(self, file_name: str | None = None) -> Tuple[str, str]:
+        """Builds paths to html and image files
+
+        Args:
+            file_name (str): file name with not extension
+
+        Returns:
+            tuple[str, str]: html file path and image file path
+        """
+
+        # check exists dirs:
+        for _dir in (self.work_dir, self.html_dir, self.images_dir):
+            if not _dir.exists():
+                _dir.mkdir(parents=True)
+
+        file_name = file_name or datetime.now().strftime("%Y_%m_%d__%H_%M")
+        html_path = self.html_dir / "{}.html".format(file_name)
+        image_path = self.images_dir / "{}.png".format(file_name)
+
+        return str(html_path), str(image_path)
+
+
+class HTMLRenderTool(MotleyTool):
+
+    def __init__(
+        self,
+        work_dir: str,
+        executable_path: str | None = None,
+        headless: bool = True,
+        window_size: Optional[Tuple[int, int]] = None,
+    ):
+        """Tool for rendering HTML as image
+
+        Args:
+            work_dir (str): Directory for saving images and html files
+        """
+        renderer = HTMLRenderer(
+            work_dir=work_dir,
+            executable_path=executable_path,
+            headless=headless,
+            window_size=window_size,
+        )
+        langchain_tool = create_render_tool(renderer)
+        super(HTMLRenderTool, self).__init__(langchain_tool)
+
+
+class HTMLRenderToolInput(BaseModel):
+    """Input for the HTMLRenderTool.
+
+    Attributes:
+        html (str):
+    """
+
+    html: str = Field(description="HTML code for rendering")
+
+
+def create_render_tool(renderer: HTMLRenderer):
+    """Create langchain tool from HTMLRenderer.render_image method
+
+    Returns:
+        Tool:
+    """
+    return Tool.from_function(
+        func=renderer.render_image,
+        name="HTML rendering tool",
+        description="A tool for rendering HTML code as an image",
+        args_schema=HTMLRenderToolInput,
+    )
diff --git a/pytest.ini b/pytest.ini
new file mode 100644
index 00000000..a27a39d2
--- /dev/null
+++ b/pytest.ini
@@ -0,0 +1,4 @@
+[pytest]
+markers =
+    fat: marks fat tests
+addopts = -m "not fat"
diff --git a/tests/test_tools/test_html_render_tool.py b/tests/test_tools/test_html_render_tool.py
new file mode 100644
index 00000000..318b9e0b
--- /dev/null
+++ b/tests/test_tools/test_html_render_tool.py
@@ -0,0 +1,25 @@
+import os
+
+import pytest
+
+from motleycrew.tools import HTMLRenderTool
+
+
+@pytest.mark.fat
+@pytest.mark.parametrize(
+    "html_code",
+    [
+        "<html><h1>Test html</h1></html>",
+        "<h1>Test html</h1>",
+    ],
+)
+def test_render_tool(tmpdir, html_code):
+    html_render_tool = HTMLRenderTool(work_dir=str(tmpdir), window_size=(800, 600), headless=False)
+
+    image_path = html_render_tool.invoke(html_code)
+    assert os.path.exists(image_path)
+    image_dir, image_file_name = os.path.split(image_path)
+    image_name = ".".join(image_file_name.split(".")[:-1])
+    html_file_name = "{}.html".format(image_name)
+    html_file_path = os.path.join(tmpdir, "html", html_file_name)
+    assert os.path.exists(html_file_path)