From 426fdafb661249476783766cbe7aff58a5012d7a Mon Sep 17 00:00:00 2001
From: H <43509927+guoyuhao2330@users.noreply.github.com>
Date: Wed, 4 Sep 2024 19:51:07 +0800
Subject: [PATCH] Add component yahoo finance  (#2244)

### What problem does this PR solve?


### Type of change

- [ ] New Feature (non-breaking change which adds functionality)
---
 agent/component/__init__.py     |  2 +
 agent/component/yahoofinance.py | 90 +++++++++++++++++++++++++++++++++
 requirements.txt                |  3 +-
 requirements_arm.txt            |  3 +-
 4 files changed, 96 insertions(+), 2 deletions(-)
 create mode 100644 agent/component/yahoofinance.py

diff --git a/agent/component/__init__.py b/agent/component/__init__.py
index c05f531f3ab..c6e5ad5782f 100644
--- a/agent/component/__init__.py
+++ b/agent/component/__init__.py
@@ -22,6 +22,8 @@
 from .baidufanyi import BaiduFanyi, BaiduFanyiParam
 from .qweather import QWeather, QWeatherParam
 from .exesql import ExeSQL, ExeSQLParam
+from .yahoofinance import YahooFinance, YahooFinanceParam
+
 
 def component_class(class_name):
     m = importlib.import_module("agent.component")
diff --git a/agent/component/yahoofinance.py b/agent/component/yahoofinance.py
new file mode 100644
index 00000000000..5bb13a42237
--- /dev/null
+++ b/agent/component/yahoofinance.py
@@ -0,0 +1,90 @@
+#
+#  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
+from agent.component.base import ComponentBase, ComponentParamBase
+import yfinance as yf
+
+
+class YahooFinanceParam(ComponentParamBase):
+    """
+    Define the YahooFinance component parameters.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.info = True
+        self.history = False
+        self.count = False
+        self.financials = False
+        self.income_stmt = False
+        self.balance_sheet = False
+        self.cash_flow_statement = False
+        self.news = True
+
+    def check(self):
+        self.check_boolean(self.info, "get all stock info")
+        self.check_boolean(self.history, "get historical market data")
+        self.check_boolean(self.count, "show share count")
+        self.check_boolean(self.financials, "show financials")
+        self.check_boolean(self.income_stmt, "income statement")
+        self.check_boolean(self.balance_sheet, "balance sheet")
+        self.check_boolean(self.cash_flow_statement, "cash flow statement")
+        self.check_boolean(self.news, "show news")
+
+
+class YahooFinance(ComponentBase, ABC):
+    component_name = "YahooFinance"
+
+    def _run(self, history, **kwargs):
+        ans = self.get_input()
+        ans = "".join(ans["content"]) if "content" in ans else ""
+        if not ans:
+            return YahooFinance.be_output("")
+
+        yohoo_res = []
+        try:
+            msft = yf.Ticker(ans)
+            if self._param.info:
+                yohoo_res.append({"content": "info:\n" + pd.Series(msft.info).to_markdown() + "\n"})
+            if self._param.history:
+                yohoo_res.append({"content": "history:\n" + msft.history().to_markdown() + "\n"})
+            if self._param.count:
+                yohoo_res.append({"content": "count:\n" + msft.get_shares_full().to_markdown() + "\n"})
+            if self._param.financials:
+                yohoo_res.append({"content": "calendar:\n" + pd.DataFrame(msft.calendar).to_markdown + "\n"})
+                yohoo_res.append({"content": "sec_filings:\n" + pd.DataFrame(msft.sec_filings).to_markdown() + "\n"})
+            if self._param.income_stmt:
+                yohoo_res.append({"content": "income statement:\n" + msft.income_stmt.to_markdown() + "\n"})
+                yohoo_res.append(
+                    {"content": "quarterly income statement:\n" + msft.quarterly_income_stmt.to_markdown() + "\n"})
+            if self._param.balance_sheet:
+                yohoo_res.append({"content": "balance sheet:\n" + msft.balance_sheet.to_markdown() + "\n"})
+                yohoo_res.append(
+                    {"content": "quarterly balance sheet:\n" + msft.quarterly_balance_sheet.to_markdown() + "\n"})
+            if self._param.cash_flow_statement:
+                yohoo_res.append({"content": "cash flow statement:\n" + msft.cashflow.to_markdown() + "\n"})
+                yohoo_res.append(
+                    {"content": "quarterly cash flow statement:\n" + msft.quarterly_cashflow.to_markdown() + "\n"})
+            if self._param.news:
+                yohoo_res.append({"content": "news:\n" + pd.DataFrame(msft.news).to_markdown() + "\n"})
+        except Exception as e:
+            print("**ERROR** " + str(e))
+
+        if not yohoo_res:
+            return YahooFinance.be_output("")
+
+        return pd.DataFrame(yohoo_res)
diff --git a/requirements.txt b/requirements.txt
index 2d089244117..f38424df610 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -94,4 +94,5 @@ wikipedia==1.4.0
 word2number==1.1
 xgboost==2.1.0
 xpinyin==0.7.6
-zhipuai==2.0.1
\ No newline at end of file
+yfinance==0.2.43
+zhipuai==2.0.1
diff --git a/requirements_arm.txt b/requirements_arm.txt
index cf0b4c504fb..8ead91b76dc 100644
--- a/requirements_arm.txt
+++ b/requirements_arm.txt
@@ -167,4 +167,5 @@ scholarly==1.7.11
 deepl==1.18.0
 psycopg2-binary==2.9.9
 tabulate==0.9.0
-vertexai==1.64.0
\ No newline at end of file
+vertexai==1.64.0
+yfinance==0.2.43