Harrison/serpapi extra tools (#691)

Co-authored-by: Bruno Bornsztein <bruno.bornsztein@gmail.com>
pull/701/head
Harrison Chase 2 years ago committed by GitHub
parent b4fcdeb56c
commit 86dbdb118b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -24,14 +24,6 @@ def _get_python_repl() -> Tool:
)
def _get_serpapi() -> Tool:
return Tool(
"Search",
SerpAPIWrapper().run,
"A search engine. Useful for when you need to answer questions about current events. Input should be a search query.",
)
def _get_requests() -> Tool:
return Tool(
"Requests",
@ -50,7 +42,6 @@ def _get_terminal() -> Tool:
_BASE_TOOLS = {
"python_repl": _get_python_repl,
"serpapi": _get_serpapi,
"requests": _get_requests,
"terminal": _get_terminal,
}
@ -139,6 +130,14 @@ def _get_google_search(**kwargs: Any) -> Tool:
)
def _get_serpapi(**kwargs: Any) -> Tool:
return Tool(
"Search",
SerpAPIWrapper(**kwargs).run,
"A search engine. Useful for when you need to answer questions about current events. Input should be a search query.",
)
_EXTRA_LLM_TOOLS = {
"news-api": (_get_news_api, ["news_api_key"]),
"tmdb-api": (_get_tmdb_api, ["tmdb_bearer_token"]),
@ -146,6 +145,7 @@ _EXTRA_LLM_TOOLS = {
_EXTRA_OPTIONAL_TOOLS = {
"wolfram-alpha": (_get_wolfram_alpha, ["wolfram_alpha_appid"]),
"google-search": (_get_google_search, ["google_api_key", "google_cse_id"]),
"serpapi": (_get_serpapi, ["serpapi_api_key"]),
}
@ -172,7 +172,7 @@ def load_tools(
elif name in _EXTRA_LLM_TOOLS:
if llm is None:
raise ValueError(f"Tool {name} requires an LLM to be provided")
_get_tool_func, extra_keys = _EXTRA_OPTIONAL_TOOLS[name]
_get_llm_tool_func, extra_keys = _EXTRA_LLM_TOOLS[name]
missing_keys = set(extra_keys).difference(kwargs)
if missing_keys:
raise ValueError(
@ -180,7 +180,7 @@ def load_tools(
f"provided: {missing_keys}"
)
sub_kwargs = {k: kwargs[k] for k in extra_keys}
tools.append(_get_tool_func(llm=llm, **sub_kwargs))
tools.append(_get_llm_tool_func(llm=llm, **sub_kwargs))
elif name in _EXTRA_OPTIONAL_TOOLS:
_get_tool_func, extra_keys = _EXTRA_OPTIONAL_TOOLS[name]
sub_kwargs = {k: kwargs[k] for k in extra_keys if k in kwargs}
@ -193,4 +193,9 @@ def load_tools(
def get_all_tool_names() -> List[str]:
"""Get a list of all possible tool names."""
return list(_BASE_TOOLS) + list(_EXTRA_OPTIONAL_TOOLS) + list(_LLM_TOOLS)
return (
list(_BASE_TOOLS)
+ list(_EXTRA_OPTIONAL_TOOLS)
+ list(_EXTRA_LLM_TOOLS)
+ list(_LLM_TOOLS)
)

Loading…
Cancel
Save