diff --git a/graph/component/__init__.py b/graph/component/__init__.py index 7bfd3f48856..222a88d96a6 100644 --- a/graph/component/__init__.py +++ b/graph/component/__init__.py @@ -10,7 +10,7 @@ from .rewrite import RewriteQuestion, RewriteQuestionParam from .keyword import KeywordExtract, KeywordExtractParam from .baidu import Baidu, BaiduParam -from .duckduckgosearch import DuckDuckGoSearch, DuckDuckGoSearchParam +from .duckduckgo import DuckDuckGo, DuckDuckGoParam def component_class(class_name): diff --git a/graph/component/duckduckgosearch.py b/graph/component/duckduckgo.py similarity index 86% rename from graph/component/duckduckgosearch.py rename to graph/component/duckduckgo.py index 8f8f20604e9..08cab52ea9b 100644 --- a/graph/component/duckduckgosearch.py +++ b/graph/component/duckduckgo.py @@ -1,62 +1,62 @@ -# -# Copyright 2024 The InfiniFlow Authors. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -import random -from abc import ABC -from functools import partial -from duckduckgo_search import DDGS -import pandas as pd - -from graph.component.base import ComponentBase, ComponentParamBase - - -class DuckDuckGoSearchParam(ComponentParamBase): - """ - Define the DuckDuckGoSearch component parameters. - """ - - def __init__(self): - super().__init__() - self.top_n = 10 - self.channel = "text" - - def check(self): - self.check_positive_integer(self.top_n, "Top N") - self.check_valid_value(self.channel, "Web Search or News", ["text", "news"]) - - -class DuckDuckGoSearch(ComponentBase, ABC): - component_name = "DuckDuckGoSearch" - - def _run(self, history, **kwargs): - ans = self.get_input() - ans = " - ".join(ans["content"]) if "content" in ans else "" - if not ans: - return DuckDuckGoSearch.be_output(self._param.no) - - if self.channel == "text": - with DDGS() as ddgs: - # {'title': '', 'href': '', 'body': ''} - duck_res = [{"content": '' + i["title"] + ' ' + i["body"]} for i in - ddgs.text(ans, max_results=self._param.top_n)] - elif self.channel == "news": - with DDGS() as ddgs: - # {'date': '', 'title': '', 'body': '', 'url': '', 'image': '', 'source': ''} - duck_res = [{"content": '' + i["title"] + ' ' + i["body"]} for i in - ddgs.news(ans, max_results=self._param.top_n)] - - df = pd.DataFrame(duck_res) - print(df, ":::::::::::::::::::::::::::::::::") - return df +# +# Copyright 2024 The InfiniFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +import random +from abc import ABC +from functools import partial +from duckduckgo_search import DDGS +import pandas as pd + +from graph.component.base import ComponentBase, ComponentParamBase + + +class DuckDuckGoParam(ComponentParamBase): + """ + Define the DuckDuckGo component parameters. + """ + + def __init__(self): + super().__init__() + self.top_n = 10 + self.channel = "text" + + def check(self): + self.check_positive_integer(self.top_n, "Top N") + self.check_valid_value(self.channel, "Web Search or News", ["text", "news"]) + + +class DuckDuckGo(ComponentBase, ABC): + component_name = "DuckDuckGo" + + def _run(self, history, **kwargs): + ans = self.get_input() + ans = " - ".join(ans["content"]) if "content" in ans else "" + if not ans: + return DuckDuckGo.be_output(self._param.no) + + if self.channel == "text": + with DDGS() as ddgs: + # {'title': '', 'href': '', 'body': ''} + duck_res = [{"content": '' + i["title"] + ' ' + i["body"]} for i in + ddgs.text(ans, max_results=self._param.top_n)] + elif self.channel == "news": + with DDGS() as ddgs: + # {'date': '', 'title': '', 'body': '', 'url': '', 'image': '', 'source': ''} + duck_res = [{"content": '' + i["title"] + ' ' + i["body"]} for i in + ddgs.news(ans, max_results=self._param.top_n)] + + df = pd.DataFrame(duck_res) + print(df, ":::::::::::::::::::::::::::::::::") + return df