H
		
	commited on
		
		
					Commit 
							
							·
						
						51856f7
	
1
								Parent(s):
							
							4ea925e
								
Add component TuShare (#2288)
Browse files### What problem does this PR solve?
### Type of change
- [x] New Feature (non-breaking change which adds functionality)
- agent/component/__init__.py +1 -0
- agent/component/tushare.py +72 -0
    	
        agent/component/__init__.py
    CHANGED
    
    | @@ -25,6 +25,7 @@ from .exesql import ExeSQL, ExeSQLParam | |
| 25 | 
             
            from .yahoofinance import YahooFinance, YahooFinanceParam
         | 
| 26 | 
             
            from .wencai import WenCai, WenCaiParam
         | 
| 27 | 
             
            from .jin10 import Jin10, Jin10Param
         | 
|  | |
| 28 |  | 
| 29 |  | 
| 30 | 
             
            def component_class(class_name):
         | 
|  | |
| 25 | 
             
            from .yahoofinance import YahooFinance, YahooFinanceParam
         | 
| 26 | 
             
            from .wencai import WenCai, WenCaiParam
         | 
| 27 | 
             
            from .jin10 import Jin10, Jin10Param
         | 
| 28 | 
            +
            from .tushare import TuShare, TuShareParam
         | 
| 29 |  | 
| 30 |  | 
| 31 | 
             
            def component_class(class_name):
         | 
    	
        agent/component/tushare.py
    ADDED
    
    | @@ -0,0 +1,72 @@ | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | 
|  | |
| 1 | 
            +
            #
         | 
| 2 | 
            +
            #  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
         | 
| 3 | 
            +
            #
         | 
| 4 | 
            +
            #  Licensed under the Apache License, Version 2.0 (the "License");
         | 
| 5 | 
            +
            #  you may not use this file except in compliance with the License.
         | 
| 6 | 
            +
            #  You may obtain a copy of the License at
         | 
| 7 | 
            +
            #
         | 
| 8 | 
            +
            #      http://www.apache.org/licenses/LICENSE-2.0
         | 
| 9 | 
            +
            #
         | 
| 10 | 
            +
            #  Unless required by applicable law or agreed to in writing, software
         | 
| 11 | 
            +
            #  distributed under the License is distributed on an "AS IS" BASIS,
         | 
| 12 | 
            +
            #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
         | 
| 13 | 
            +
            #  See the License for the specific language governing permissions and
         | 
| 14 | 
            +
            #  limitations under the License.
         | 
| 15 | 
            +
            #
         | 
| 16 | 
            +
            import json
         | 
| 17 | 
            +
            from abc import ABC
         | 
| 18 | 
            +
            import pandas as pd
         | 
| 19 | 
            +
            import time
         | 
| 20 | 
            +
            import requests
         | 
| 21 | 
            +
            from agent.component.base import ComponentBase, ComponentParamBase
         | 
| 22 | 
            +
             | 
| 23 | 
            +
             | 
| 24 | 
            +
            class TuShareParam(ComponentParamBase):
         | 
| 25 | 
            +
                """
         | 
| 26 | 
            +
                Define the TuShare component parameters.
         | 
| 27 | 
            +
                """
         | 
| 28 | 
            +
             | 
| 29 | 
            +
                def __init__(self):
         | 
| 30 | 
            +
                    super().__init__()
         | 
| 31 | 
            +
                    self.token = "xxx"
         | 
| 32 | 
            +
                    self.src = "eastmoney"
         | 
| 33 | 
            +
                    self.start_date = "2024-01-01 09:00:00"
         | 
| 34 | 
            +
                    self.end_date = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
         | 
| 35 | 
            +
                    self.keyword = ""
         | 
| 36 | 
            +
             | 
| 37 | 
            +
                def check(self):
         | 
| 38 | 
            +
                    self.check_valid_value(self.src, "Quick News Source",
         | 
| 39 | 
            +
                                           ["sina", "wallstreetcn", "10jqka", "eastmoney", "yuncaijing", "fenghuang", "jinrongjie"])
         | 
| 40 | 
            +
             | 
| 41 | 
            +
             | 
| 42 | 
            +
            class TuShare(ComponentBase, ABC):
         | 
| 43 | 
            +
                component_name = "TuShare"
         | 
| 44 | 
            +
             | 
| 45 | 
            +
                def _run(self, history, **kwargs):
         | 
| 46 | 
            +
                    ans = self.get_input()
         | 
| 47 | 
            +
                    ans = ",".join(ans["content"]) if "content" in ans else ""
         | 
| 48 | 
            +
                    if not ans:
         | 
| 49 | 
            +
                        return TuShare.be_output("")
         | 
| 50 | 
            +
             | 
| 51 | 
            +
                    try:
         | 
| 52 | 
            +
                        tus_res = []
         | 
| 53 | 
            +
                        params = {
         | 
| 54 | 
            +
                            "api_name": "news",
         | 
| 55 | 
            +
                            "token": self._param.token,
         | 
| 56 | 
            +
                            "params": {"src": self._param.src, "start_date": self._param.start_date,
         | 
| 57 | 
            +
                                       "end_date": self._param.end_date}
         | 
| 58 | 
            +
                        }
         | 
| 59 | 
            +
                        response = requests.post(url="http://api.tushare.pro", data=json.dumps(params).encode('utf-8'))
         | 
| 60 | 
            +
                        response = response.json()
         | 
| 61 | 
            +
                        if response['code'] != 0:
         | 
| 62 | 
            +
                            return TuShare.be_output(response['msg'])
         | 
| 63 | 
            +
                        df = pd.DataFrame(response['data']['items'])
         | 
| 64 | 
            +
                        df.columns = response['data']['fields']
         | 
| 65 | 
            +
                        tus_res.append({"content": (df[df['content'].str.contains(self._param.keyword, case=False)]).to_markdown()})
         | 
| 66 | 
            +
                    except Exception as e:
         | 
| 67 | 
            +
                        return TuShare.be_output("**ERROR**: " + str(e))
         | 
| 68 | 
            +
             | 
| 69 | 
            +
                    if not tus_res:
         | 
| 70 | 
            +
                        return TuShare.be_output("")
         | 
| 71 | 
            +
             | 
| 72 | 
            +
                    return pd.DataFrame(tus_res)
         |