Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Html rendering tool #59

Merged
merged 6 commits into from
Jul 4, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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)
Loading