diff --git a/.env.template b/.env.template index 156aaea..423e91c 100644 --- a/.env.template +++ b/.env.template @@ -2,4 +2,5 @@ OPENAI_API_KEY="" SEARCHAPI_API_KEY="" GOOGLE_CSE_ID="" GOOGLE_API_KEY="" -BING_API_KEY="" \ No newline at end of file +BING_API_KEY="" +BOCHA_API_KEY="" \ No newline at end of file diff --git a/swarm/environment/__init__.py b/swarm/environment/__init__.py index 35e9d53..cd2edba 100644 --- a/swarm/environment/__init__.py +++ b/swarm/environment/__init__.py @@ -1,10 +1,11 @@ from swarm.environment.tools.reader.readers import GeneralReader -from swarm.environment.tools.search.search import GoogleSearchEngine, SearchAPIEngine, BingSearchEngine +from swarm.environment.tools.search.search import GoogleSearchEngine, SearchAPIEngine, BingSearchEngine,BoChaSearchEngine __all__ = [ "GeneralReader", "GoogleSearchEngine", "SearchAPIEngine", - "BingSearchEngine" + "BingSearchEngine", + "BoChaSearchEngine" ] \ No newline at end of file diff --git a/swarm/environment/operations/web_search.py b/swarm/environment/operations/web_search.py index 504e9f5..3b30196 100644 --- a/swarm/environment/operations/web_search.py +++ b/swarm/environment/operations/web_search.py @@ -9,7 +9,7 @@ from swarm.llm.format import Message from swarm.graph import Node -from swarm.environment import GoogleSearchEngine, SearchAPIEngine, BingSearchEngine +from swarm.environment import GoogleSearchEngine, SearchAPIEngine, BingSearchEngine, BoChaSearchEngine from swarm.utils.log import logger, swarmlog from swarm.utils.globals import Cost from swarm.environment.prompt.prompt_set_registry import PromptSetRegistry @@ -41,6 +41,8 @@ def _get_searcher(self): return SearchAPIEngine() if os.getenv("GOOGLE_API_KEY"): return GoogleSearchEngine() + if os.getenv("BOCHA_API_KEY"): + return BoChaSearchEngine() async def _execute(self, inputs: List[Any] = [], max_keywords: int = 5, **kwargs): diff --git a/swarm/environment/tools/search/search.py b/swarm/environment/tools/search/search.py index 9aaa31f..a6f4a13 100644 --- a/swarm/environment/tools/search/search.py +++ b/swarm/environment/tools/search/search.py @@ -6,6 +6,31 @@ import requests import ast + +class BoChaSearchEngine(): + def __init__(self) -> None: + self.api_key = os.getenv("BOCHA_API_KEY") + + def search(self, query: str, num: int = 3): + try: + url = 'https://api.bochaai.com/v1/web-search' + headers = { + 'Authorization': f'Bearer {self.api_key}', # 请替换为你的API密钥 + 'Content-Type': 'application/json' + } + data = { + "query": query, + } + + response = requests.post(url, headers=headers, json=data) + + json_response = response.json() + res = json_response["data"]["webPages"]["value"] + print(res) + return '\n'.join([item['snippet'] for item in res]) + except: + return 'Cannot get search results from BoCha API' + class BingSearchEngine(): def __init__(self) -> None: self.api_key = os.getenv("BING_API_KEY")