Adding new description generator

This commit is contained in:
João Moura
2024-02-29 03:09:48 -03:00
parent 640b5a9461
commit ec97e15a3a
19 changed files with 53 additions and 10 deletions

View File

@@ -13,10 +13,12 @@ class BaseTool(BaseModel, ABC):
"""Used to tell the model how/when/why to use the tool."""
args_schema: Optional[Type[V1BaseModel]] = None
"""The schema for the arguments that the tool accepts."""
description_updated: bool = False
@model_validator(mode="after")
def _check_args_schema(self):
self._set_args_schema()
self._generate_description()
return self
def run(
@@ -56,6 +58,13 @@ class BaseTool(BaseModel, ABC):
},
},
)
def _generate_description(self):
args = []
for arg, attribute in self.args_schema.schema()['properties'].items():
args.append(f"{arg}: '{attribute['type']}'")
description = self.description.replace('\n', ' ')
self.description = f"{self.name}({', '.join(args)}) - {description}"
class Tool(BaseTool):

View File

@@ -28,6 +28,7 @@ class CodeDocsSearchTool(RagTool):
self.docs_url = docs_url
self.description = f"A tool that can be used to semantic search a query the {docs_url} Code Docs content."
self.args_schema = FixedCodeDocsSearchToolSchema
self._generate_description()
def _run(
self,

View File

@@ -28,6 +28,7 @@ class CSVSearchTool(RagTool):
self.csv = csv
self.description = f"A tool that can be used to semantic search a query the {csv} CSV's content."
self.args_schema = FixedCSVSearchToolSchema
self._generate_description()
def _run(
self,

View File

@@ -23,6 +23,7 @@ class DirectoryReadTool(BaseTool):
self.directory = directory
self.description = f"A tool that can be used to list {directory}'s content."
self.args_schema = FixedDirectoryReadToolSchema
self._generate_description()
def _run(
self,

View File

@@ -28,6 +28,7 @@ class DirectorySearchTool(RagTool):
self.directory = directory
self.description = f"A tool that can be used to semantic search a query the {directory} directory's content."
self.args_schema = FixedDirectorySearchToolSchema
self._generate_description()
def _run(
self,

View File

@@ -28,6 +28,7 @@ class DOCXSearchTool(RagTool):
self.docx = docx
self.description = f"A tool that can be used to semantic search a query the {docx} DOCX's content."
self.args_schema = FixedDOCXSearchToolSchema
self._generate_description()
def _run(
self,

View File

@@ -22,6 +22,7 @@ class FileReadTool(BaseTool):
self.file_path = file_path
self.description = f"A tool that can be used to read {file_path}'s content."
self.args_schema = FixedFileReadToolSchema
self._generate_description()
def _run(
self,

View File

@@ -31,6 +31,7 @@ class GithubSearchTool(RagTool):
self.github_repo = github_repo
self.description = f"A tool that can be used to semantic search a query the {github_repo} github repo's content."
self.args_schema = FixedGithubSearchToolSchema
self._generate_description()
def _run(
self,

View File

@@ -28,6 +28,7 @@ class JSONSearchTool(RagTool):
self.json_path = json_path
self.description = f"A tool that can be used to semantic search a query the {json} JSON's content."
self.args_schema = FixedJSONSearchToolSchema
self._generate_description()
def _run(
self,

View File

@@ -28,6 +28,7 @@ class MDXSearchTool(RagTool):
self.mdx = mdx
self.description = f"A tool that can be used to semantic search a query the {mdx} MDX's content."
self.args_schema = FixedMDXSearchToolSchema
self._generate_description()
def _run(
self,

View File

@@ -28,6 +28,7 @@ class PDFSearchTool(RagTool):
self.pdf = pdf
self.description = f"A tool that can be used to semantic search a query the {pdf} PDF's content."
self.args_schema = FixedPDFSearchToolSchema
self._generate_description()
def _run(
self,

View File

@@ -24,6 +24,7 @@ class PGSearchTool(RagTool):
if table_name is not None:
self.table_name = table_name
self.description = f"A tool that can be used to semantic search a query the {table_name} database table's content."
self._generate_description()
else:
raise('To use PGSearchTool, you must provide a `table_name` argument')

View File

@@ -1,3 +1,4 @@
import os
import requests
from bs4 import BeautifulSoup
from typing import Optional, Type, Any
@@ -18,20 +19,28 @@ class ScrapeElementFromWebsiteTool(BaseTool):
description: str = "A tool that can be used to read a website content."
args_schema: Type[BaseModel] = ScrapeElementFromWebsiteToolSchema
website_url: Optional[str] = None
cookies: Optional[dict] = None
css_element: Optional[str] = None
headers: Optional[dict] = {
'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3',
'Accept-Language': 'en-US,en;q=0.5',
'Referer': 'https://www.google.com/'
'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/96.0.4664.110 Safari/537.36',
'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.9',
'Accept-Language': 'en-US,en;q=0.9',
'Referer': 'https://www.google.com/',
'Connection': 'keep-alive',
'Upgrade-Insecure-Requests': '1',
'Accept-Encoding': 'gzip, deflate, br'
}
def __init__(self, website_url: Optional[str] = None, css_element: Optional[str] = None, **kwargs):
def __init__(self, website_url: Optional[str] = None, cookies: Optional[dict] = None, css_element: Optional[str] = None, **kwargs):
super().__init__(**kwargs)
if website_url is not None:
self.website_url = website_url
self.css_element = css_element
self.description = f"A tool that can be used to read {website_url}'s content."
self.args_schema = FixedScrapeElementFromWebsiteToolSchema
self._generate_description()
if cookies is not None:
self.cookies = {cookies["name"]: os.getenv(cookies["value"])}
def _run(
self,
@@ -39,7 +48,7 @@ class ScrapeElementFromWebsiteTool(BaseTool):
) -> Any:
website_url = kwargs.get('website_url', self.website_url)
css_element = kwargs.get('css_element', self.css_element)
page = requests.get(website_url, headers=self.headers)
page = requests.get(website_url, headers=self.headers, cookies=self.cookies if self.cookies else {})
parsed = BeautifulSoup(page.content, "html.parser")
elements = parsed.select(css_element)
return "\n".join([element.get_text() for element in elements])

View File

@@ -1,3 +1,4 @@
import os
import requests
from bs4 import BeautifulSoup
from typing import Optional, Type, Any
@@ -17,25 +18,33 @@ class ScrapeWebsiteTool(BaseTool):
description: str = "A tool that can be used to read a website content."
args_schema: Type[BaseModel] = ScrapeWebsiteToolSchema
website_url: Optional[str] = None
cookies: Optional[dict] = None
headers: Optional[dict] = {
'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3',
'Accept-Language': 'en-US,en;q=0.5',
'Referer': 'https://www.google.com/'
'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/96.0.4664.110 Safari/537.36',
'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.9',
'Accept-Language': 'en-US,en;q=0.9',
'Referer': 'https://www.google.com/',
'Connection': 'keep-alive',
'Upgrade-Insecure-Requests': '1',
'Accept-Encoding': 'gzip, deflate, br'
}
def __init__(self, website_url: Optional[str] = None, **kwargs):
def __init__(self, website_url: Optional[str] = None, cookies: Optional[dict] = None, **kwargs):
super().__init__(**kwargs)
if website_url is not None:
self.website_url = website_url
self.description = f"A tool that can be used to read {website_url}'s content."
self.args_schema = FixedScrapeWebsiteToolSchema
self._generate_description()
if cookies is not None:
self.cookies = {cookies["name"]: os.getenv(cookies["value"])}
def _run(
self,
**kwargs: Any,
) -> Any:
website_url = kwargs.get('website_url', self.website_url)
page = requests.get(website_url, headers=self.headers)
page = requests.get(website_url, headers=self.headers, cookies=self.cookies if self.cookies else {})
parsed = BeautifulSoup(page.content, "html.parser")
return parsed.get_text()

View File

@@ -27,6 +27,7 @@ class TXTSearchTool(RagTool):
self.txt = txt
self.description = f"A tool that can be used to semantic search a query the {txt} txt's content."
self.args_schema = FixedTXTSearchToolSchema
self._generate_description()
def _run(
self,

View File

@@ -28,6 +28,7 @@ class WebsiteSearchTool(RagTool):
self.website = website
self.description = f"A tool that can be used to semantic search a query from {website} website content."
self.args_schema = FixedWebsiteSearchToolSchema
self._generate_description()
def _run(
self,

View File

@@ -28,6 +28,7 @@ class XMLSearchTool(RagTool):
self.xml = xml
self.description = f"A tool that can be used to semantic search a query the {xml} XML's content."
self.args_schema = FixedXMLSearchToolSchema
self._generate_description()
def _run(
self,

View File

@@ -28,6 +28,7 @@ class YoutubeChannelSearchTool(RagTool):
self.youtube_channel_handle = youtube_channel_handle
self.description = f"A tool that can be used to semantic search a query the {youtube_channel_handle} Youtube Channels content."
self.args_schema = FixedYoutubeChannelSearchToolSchema
self._generate_description()
def _run(
self,

View File

@@ -28,6 +28,7 @@ class YoutubeVideoSearchTool(RagTool):
self.youtube_video_url = youtube_video_url
self.description = f"A tool that can be used to semantic search a query the {youtube_video_url} Youtube Video content."
self.args_schema = FixedYoutubeVideoSearchToolSchema
self._generate_description()
def _run(
self,