-
Notifications
You must be signed in to change notification settings - Fork 135
/
tool.py
82 lines (70 loc) · 2.63 KB
/
tool.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
from __future__ import annotations
import logging
import json
from attr import define, field
from griptape.artifacts import BaseArtifact, TextArtifact, ErrorArtifact, ListArtifact
from griptape.loaders import TextLoader
from schema import Schema, Literal
from griptape.tools import BaseTool
from griptape.utils.decorators import activity
@define
class WebScraper(BaseTool):
include_links: bool = field(default=True, kw_only=True)
@activity(config={
"description": "Can be used to browse a web page and load its content",
"schema": Schema({
Literal(
"url",
description="Valid HTTP URL"
): str
})
})
def get_content(self, params: dict) -> ListArtifact | ErrorArtifact:
url = params["values"]["url"]
page = self._load_page(url)
if isinstance(page, ErrorArtifact):
return page
else:
return ListArtifact(
TextLoader().text_to_artifacts(page.get("text"))
)
@activity(config={
"description": "Can be used to load a web page author",
"schema": Schema({
Literal(
"url",
description="Valid HTTP URL"
): str
})
})
def get_author(self, params: dict) -> BaseArtifact:
url = params["values"]["url"]
page = self._load_page(url)
if isinstance(page, ErrorArtifact):
return page
else:
return TextArtifact(page.get("author"))
def _load_page(self, url: str) -> dict | ErrorArtifact:
import trafilatura
from trafilatura.settings import use_config
config = use_config()
page = trafilatura.fetch_url(url, no_ssl=True)
# This disables signal, so that trafilatura can work on any thread:
# More info: https://trafilatura.readthedocs.io/usage-python.html#disabling-signal
config.set("DEFAULT", "EXTRACTION_TIMEOUT", "0")
# Disable error logging in trafilatura as it sometimes logs errors from lxml, even though
# the end result of page parsing is successful.
logging.getLogger("trafilatura").setLevel(logging.FATAL)
if page is None:
return ErrorArtifact("error: can't access URL")
else:
content = trafilatura.extract(
page,
include_links=self.include_links,
output_format="json",
config=config
)
if content:
return json.loads(content)
else:
return ErrorArtifact("error: can't load web page content")