Spaces:
Sleeping
Sleeping
import os | |
from ..tools.tool import Tool, ToolSchema | |
from ..tools.web_search_utils import get_websearcher_cls | |
from ..tools.web_search_utils.search_util import \ | |
AuthenticationKey | |
from pydantic import ValidationError | |
class WebSearch(Tool): | |
description = 'surfacing relevant information from billions of web documents. Help ' \ | |
'you find what you are looking for from the world-wide-web to comb ' \ | |
'billions of webpages, images, videos, and news.' | |
name = 'web_search_utils' | |
parameters: list = [{ | |
'name': 'query', | |
'description': | |
"""The user's search query term. The term may not be empty.""", | |
'required': True | |
}] | |
def __init__(self, cfg={}): | |
super().__init__() | |
available_searchers = get_websearcher_cls() | |
all_searchers = AuthenticationKey.to_dict() | |
if not len(available_searchers): | |
raise ValueError( | |
f'At least one of web search api token should be set: {all_searchers}' | |
) | |
searcher = cfg.pop('searcher', None) | |
if not searcher: | |
self.searcher = available_searchers[0](**cfg) | |
else: | |
if isinstance(searcher, | |
str) and len(searcher) and all_searchers.get( | |
searcher, None): | |
cls = available_searchers.get(searcher, None) | |
if not cls: | |
raise ValueError( | |
f'The searcher {searcher}\'s token is not set: {all_searchers.get(searcher, None)}' | |
) | |
self.searcher = cls(**cfg) | |
else: | |
raise ValueError( | |
f'The searcher {searcher} should be one of {all_searchers.keys()}' | |
) | |
try: | |
all_para = { | |
'name': self.name, | |
'description': self.description, | |
'parameters': self.parameters | |
} | |
self.tool_schema = ToolSchema(**all_para) | |
except ValidationError: | |
raise ValueError(f'Error when parsing parameters of {self.name}') | |
self.is_remote_tool = True | |
self._str = self.tool_schema.model_dump_json() | |
self._function = self.parse_pydantic_model_to_openai_function(all_para) | |
def _remote_call(self, *args, **kwargs): | |
query = self._handle_input_fallback(**kwargs) | |
if not query or not len(query): | |
raise ValueError( | |
'parameter `query` of tool web-search is None or Empty.') | |
res = self.searcher(query) | |
return {'result': [item.__dict__ for item in res]} | |
def _handle_input_fallback(self, **kwargs): | |
query = kwargs.get('query', None) | |
fallback = kwargs.get('fallback', None) | |
if query and isinstance(query, str) and len(query): | |
return query | |
else: | |
return fallback | |
if __name__ == '__main__': | |
tool = WebSearch() | |
res = tool(query='2024年 元旦 哈尔滨天气') | |
print(res) | |