diff --git a/agent/component/__init__.py b/agent/component/__init__.py index 778189c3..2a24f229 100644 --- a/agent/component/__init__.py +++ b/agent/component/__init__.py @@ -26,6 +26,7 @@ from .yahoofinance import YahooFinance, YahooFinanceParam from .wencai import WenCai, WenCaiParam from .jin10 import Jin10, Jin10Param from .tushare import TuShare, TuShareParam +from .akshare import AkShare, AkShareParam def component_class(class_name): diff --git a/agent/component/akshare.py b/agent/component/akshare.py new file mode 100644 index 00000000..c6ed2fb6 --- /dev/null +++ b/agent/component/akshare.py @@ -0,0 +1,56 @@ +# +# 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. +# +from abc import ABC +import pandas as pd +import akshare as ak +from agent.component.base import ComponentBase, ComponentParamBase + + +class AkShareParam(ComponentParamBase): + """ + Define the AkShare component parameters. + """ + + def __init__(self): + super().__init__() + self.top_n = 10 + + def check(self): + self.check_positive_integer(self.top_n, "Top N") + + +class AkShare(ComponentBase, ABC): + component_name = "AkShare" + + def _run(self, history, **kwargs): + ans = self.get_input() + ans = ",".join(ans["content"]) if "content" in ans else "" + if not ans: + return AkShare.be_output("") + + try: + ak_res = [] + stock_news_em_df = ak.stock_news_em(symbol=ans) + stock_news_em_df = stock_news_em_df.head(self._param.top_n) + ak_res = [{"content": '' + i["新闻标题"] + '\n 新闻内容: ' + i[ + "新闻内容"] + " \n发布时间:" + i["发布时间"] + " \n文章来源: " + i["文章来源"]} for index, i in stock_news_em_df.iterrows()] + except Exception as e: + return AkShare.be_output("**ERROR**: " + str(e)) + + if not ak_res: + return AkShare.be_output("") + + return pd.DataFrame(ak_res) diff --git a/requirements.txt b/requirements.txt index b1da26c5..9f86aaeb 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,5 @@ -anthropic===0.34.1 +akshare==1.14.72 +anthropic==0.34.1 arxiv==2.1.3 Aspose.Slides==24.2.0 BCEmbedding==0.1.3 diff --git a/requirements_arm.txt b/requirements_arm.txt index 3b73a2ce..87540366 100644 --- a/requirements_arm.txt +++ b/requirements_arm.txt @@ -170,3 +170,4 @@ tabulate==0.9.0 vertexai==1.64.0 yfinance==0.2.43 pywencai==0.12.2 +akshare==1.14.72