Merge pull request #107 from masorensen/bug/ragTools-kwargs

Scoping issue causing error in RAG tools
This commit is contained in:
João Moura
2024-10-17 16:43:11 -03:00
committed by GitHub
15 changed files with 19 additions and 19 deletions

View File

@@ -31,6 +31,7 @@ class CodeDocsSearchTool(RagTool):
def __init__(self, docs_url: Optional[str] = None, **kwargs):
super().__init__(**kwargs)
if docs_url is not None:
kwargs["data_type"] = DataType.DOCS_SITE
self.add(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
@@ -41,7 +42,6 @@ class CodeDocsSearchTool(RagTool):
*args: Any,
**kwargs: Any,
) -> None:
kwargs["data_type"] = DataType.DOCS_SITE
super().add(*args, **kwargs)
def _before_run(

View File

@@ -31,6 +31,7 @@ class CSVSearchTool(RagTool):
def __init__(self, csv: Optional[str] = None, **kwargs):
super().__init__(**kwargs)
if csv is not None:
kwargs["data_type"] = DataType.CSV
self.add(csv)
self.description = f"A tool that can be used to semantic search a query the {csv} CSV's content."
self.args_schema = FixedCSVSearchToolSchema
@@ -41,7 +42,6 @@ class CSVSearchTool(RagTool):
*args: Any,
**kwargs: Any,
) -> None:
kwargs["data_type"] = DataType.CSV
super().add(*args, **kwargs)
def _before_run(

View File

@@ -31,6 +31,7 @@ class DirectorySearchTool(RagTool):
def __init__(self, directory: Optional[str] = None, **kwargs):
super().__init__(**kwargs)
if directory is not None:
kwargs["loader"] = DirectoryLoader(config=dict(recursive=True))
self.add(directory)
self.description = f"A tool that can be used to semantic search a query the {directory} directory's content."
self.args_schema = FixedDirectorySearchToolSchema
@@ -41,7 +42,6 @@ class DirectorySearchTool(RagTool):
*args: Any,
**kwargs: Any,
) -> None:
kwargs["loader"] = DirectoryLoader(config=dict(recursive=True))
super().add(*args, **kwargs)
def _before_run(

View File

@@ -37,6 +37,7 @@ class DOCXSearchTool(RagTool):
def __init__(self, docx: Optional[str] = None, **kwargs):
super().__init__(**kwargs)
if docx is not None:
kwargs["data_type"] = DataType.DOCX
self.add(docx)
self.description = f"A tool that can be used to semantic search a query the {docx} DOCX's content."
self.args_schema = FixedDOCXSearchToolSchema
@@ -47,7 +48,6 @@ class DOCXSearchTool(RagTool):
*args: Any,
**kwargs: Any,
) -> None:
kwargs["data_type"] = DataType.DOCX
super().add(*args, **kwargs)
def _before_run(

View File

@@ -38,6 +38,9 @@ class GithubSearchTool(RagTool):
def __init__(self, github_repo: Optional[str] = None, **kwargs):
super().__init__(**kwargs)
if github_repo is not None:
kwargs["data_type"] = "github"
kwargs["loader"] = GithubLoader(config={"token": self.gh_token})
self.add(repo=github_repo)
self.description = f"A tool that can be used to semantic search a query the {github_repo} github repo's content. This is not the GitHub API, but instead a tool that can provide semantic search capabilities."
self.args_schema = FixedGithubSearchToolSchema
@@ -51,8 +54,6 @@ class GithubSearchTool(RagTool):
) -> None:
content_types = content_types or self.content_types
kwargs["data_type"] = "github"
kwargs["loader"] = GithubLoader(config={"token": self.gh_token})
super().add(f"repo:{repo} type:{','.join(content_types)}", **kwargs)
def _before_run(

View File

@@ -31,6 +31,7 @@ class JSONSearchTool(RagTool):
def __init__(self, json_path: Optional[str] = None, **kwargs):
super().__init__(**kwargs)
if json_path is not None:
kwargs["data_type"] = DataType.JSON
self.add(json_path)
self.description = f"A tool that can be used to semantic search a query the {json_path} JSON's content."
self.args_schema = FixedJSONSearchToolSchema
@@ -41,7 +42,6 @@ class JSONSearchTool(RagTool):
*args: Any,
**kwargs: Any,
) -> None:
kwargs["data_type"] = DataType.JSON
super().add(*args, **kwargs)
def _before_run(

View File

@@ -31,6 +31,7 @@ class MDXSearchTool(RagTool):
def __init__(self, mdx: Optional[str] = None, **kwargs):
super().__init__(**kwargs)
if mdx is not None:
kwargs["data_type"] = DataType.MDX
self.add(mdx)
self.description = f"A tool that can be used to semantic search a query the {mdx} MDX's content."
self.args_schema = FixedMDXSearchToolSchema
@@ -41,7 +42,6 @@ class MDXSearchTool(RagTool):
*args: Any,
**kwargs: Any,
) -> None:
kwargs["data_type"] = DataType.MDX
super().add(*args, **kwargs)
def _before_run(

View File

@@ -25,6 +25,8 @@ class MySQLSearchTool(RagTool):
def __init__(self, table_name: str, **kwargs):
super().__init__(**kwargs)
kwargs["data_type"] = "mysql"
kwargs["loader"] = MySQLLoader(config=dict(url=self.db_uri))
self.add(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()
@@ -34,8 +36,6 @@ class MySQLSearchTool(RagTool):
table_name: str,
**kwargs: Any,
) -> None:
kwargs["data_type"] = "mysql"
kwargs["loader"] = MySQLLoader(config=dict(url=self.db_uri))
super().add(f"SELECT * FROM {table_name};", **kwargs)
def _run(

View File

@@ -30,6 +30,7 @@ class PDFSearchTool(RagTool):
def __init__(self, pdf: Optional[str] = None, **kwargs):
super().__init__(**kwargs)
if pdf is not None:
kwargs["data_type"] = DataType.PDF_FILE
self.add(pdf)
self.description = f"A tool that can be used to semantic search a query the {pdf} PDF's content."
self.args_schema = FixedPDFSearchToolSchema
@@ -56,7 +57,6 @@ class PDFSearchTool(RagTool):
*args: Any,
**kwargs: Any,
) -> None:
kwargs["data_type"] = DataType.PDF_FILE
super().add(*args, **kwargs)
def _before_run(

View File

@@ -25,6 +25,8 @@ class PGSearchTool(RagTool):
def __init__(self, table_name: str, **kwargs):
super().__init__(**kwargs)
kwargs["data_type"] = "postgres"
kwargs["loader"] = PostgresLoader(config=dict(url=self.db_uri))
self.add(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()
@@ -34,8 +36,6 @@ class PGSearchTool(RagTool):
table_name: str,
**kwargs: Any,
) -> None:
kwargs["data_type"] = "postgres"
kwargs["loader"] = PostgresLoader(config=dict(url=self.db_uri))
super().add(f"SELECT * FROM {table_name};", **kwargs)
def _run(

View File

@@ -31,6 +31,7 @@ class TXTSearchTool(RagTool):
def __init__(self, txt: Optional[str] = None, **kwargs):
super().__init__(**kwargs)
if txt is not None:
kwargs["data_type"] = DataType.TEXT_FILE
self.add(txt)
self.description = f"A tool that can be used to semantic search a query the {txt} txt's content."
self.args_schema = FixedTXTSearchToolSchema
@@ -41,7 +42,6 @@ class TXTSearchTool(RagTool):
*args: Any,
**kwargs: Any,
) -> None:
kwargs["data_type"] = DataType.TEXT_FILE
super().add(*args, **kwargs)
def _before_run(

View File

@@ -33,6 +33,7 @@ class WebsiteSearchTool(RagTool):
def __init__(self, website: Optional[str] = None, **kwargs):
super().__init__(**kwargs)
if website is not None:
kwargs["data_type"] = DataType.WEB_PAGE
self.add(website)
self.description = f"A tool that can be used to semantic search a query from {website} website content."
self.args_schema = FixedWebsiteSearchToolSchema
@@ -43,7 +44,6 @@ class WebsiteSearchTool(RagTool):
*args: Any,
**kwargs: Any,
) -> None:
kwargs["data_type"] = DataType.WEB_PAGE
super().add(*args, **kwargs)
def _before_run(

View File

@@ -31,6 +31,7 @@ class XMLSearchTool(RagTool):
def __init__(self, xml: Optional[str] = None, **kwargs):
super().__init__(**kwargs)
if xml is not None:
kwargs["data_type"] = DataType.XML
self.add(xml)
self.description = f"A tool that can be used to semantic search a query the {xml} XML's content."
self.args_schema = FixedXMLSearchToolSchema
@@ -41,7 +42,6 @@ class XMLSearchTool(RagTool):
*args: Any,
**kwargs: Any,
) -> None:
kwargs["data_type"] = DataType.XML
super().add(*args, **kwargs)
def _before_run(

View File

@@ -33,6 +33,7 @@ class YoutubeChannelSearchTool(RagTool):
def __init__(self, youtube_channel_handle: Optional[str] = None, **kwargs):
super().__init__(**kwargs)
if youtube_channel_handle is not None:
kwargs["data_type"] = DataType.YOUTUBE_CHANNEL
self.add(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
@@ -45,8 +46,6 @@ class YoutubeChannelSearchTool(RagTool):
) -> None:
if not youtube_channel_handle.startswith("@"):
youtube_channel_handle = f"@{youtube_channel_handle}"
kwargs["data_type"] = DataType.YOUTUBE_CHANNEL
super().add(youtube_channel_handle, **kwargs)
def _before_run(

View File

@@ -33,6 +33,7 @@ class YoutubeVideoSearchTool(RagTool):
def __init__(self, youtube_video_url: Optional[str] = None, **kwargs):
super().__init__(**kwargs)
if youtube_video_url is not None:
kwargs["data_type"] = DataType.YOUTUBE_VIDEO
self.add(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
@@ -43,7 +44,6 @@ class YoutubeVideoSearchTool(RagTool):
*args: Any,
**kwargs: Any,
) -> None:
kwargs["data_type"] = DataType.YOUTUBE_VIDEO
super().add(*args, **kwargs)
def _before_run(