Skip to content

Commit

Permalink
Html rendering tool (#59)
Browse files Browse the repository at this point in the history
* add HtmlRenderTool

* remove HtmlRenderTool from tools init

* add HtmlRenderTool to tools init

* add HtmlRenderTool tests

* update PythonLinterTool

* minor improvements

---------

Co-authored-by: User <[email protected]>
Co-authored-by: whimo <[email protected]>
  • Loading branch information
3 people committed Jul 4, 2024
1 parent e4b9fd3 commit 1938a62
Show file tree
Hide file tree
Showing 6 changed files with 207 additions and 16 deletions.
30 changes: 19 additions & 11 deletions motleycrew/common/exceptions.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -21,7 +23,7 @@ def __str__(self) -> str:

class LLMFrameworkNotSupported(Exception):
def __init__(self, llm_framework: str):
""" Description
"""Description
Args:
llm_framework (str):
Expand All @@ -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

Expand All @@ -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

Expand All @@ -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
Expand All @@ -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):
Expand All @@ -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)

Expand All @@ -133,4 +140,5 @@ def __str__(self):

class InvalidOutput(Exception):
"""Raised in output handlers when an agent's output is not accepted"""

pass
7 changes: 4 additions & 3 deletions motleycrew/tools/__init__.py
Original file line number Diff line number Diff line change
@@ -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
7 changes: 5 additions & 2 deletions motleycrew/tools/code/python_linter.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):

Expand Down
150 changes: 150 additions & 0 deletions motleycrew/tools/html_render_tool.py
Original file line number Diff line number Diff line change
@@ -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,
)
4 changes: 4 additions & 0 deletions pytest.ini
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
[pytest]
markers =
fat: marks fat tests
addopts = -m "not fat"
25 changes: 25 additions & 0 deletions tests/test_tools/test_html_render_tool.py
Original file line number Diff line number Diff line change
@@ -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)

0 comments on commit 1938a62

Please sign in to comment.