tools.py

from smolagents import Tool
import os
import requests
import json
from dotenv import load_dotenv
import tushare as ts

load_dotenv()

_pro_api = None


def get_tushare_pro():
    global _pro_api
    if _pro_api is None:
        token = os.getenv("TUSHARE_TOKEN")
        if not token:
            raise RuntimeError("环境变量 TUSHARE_TOKEN 未设置，无法调用 Tushare 接口")
        ts.set_token(token)
        _pro_api = ts.pro_api()
    return _pro_api


class TushareDailyDataFetcher(Tool):
    name = "tushare_daily_data_fetcher"
    description = """
    用于根据股票代码和日期范围获取Tushare的每日行情数据（沪深京A股日频行情），
    数据会保存为CSV文件并返回该文件的绝对路径"""
    inputs = {
        "stock_code": {
            "type": "string",
            "description": "股票代码，需带交易所后缀，如 '000001.SZ'、'600519.SH'",
        },
        "start_date": {
            "type": "string",
            "description": "开始日期，格式 YYYYMMDD，如 '20250101'",
            "nullable": True,
        },
        "end_date": {
            "type": "string",
            "description": "结束日期，格式 YYYYMMDD，如 '20251231'",
            "nullable": True,
        },
    }
    output_type = "string"

    def forward(
        self, stock_code: str, start_date: str = None, end_date: str = None
    ) -> str:
        pro = get_tushare_pro()
        df = pro.daily(ts_code=stock_code, start_date=start_date, end_date=end_date)
        if df is None or df.empty:
            return f"未获取到 {stock_code} 的日线行情数据，请检查股票代码和日期范围"
        # Tushare 默认按日期倒序，改为升序便于阅读和分析
        df = df.sort_values("trade_date").reset_index(drop=True)
        os.makedirs("data", exist_ok=True)
        filename = f"data/{stock_code}_{start_date or 'all'}_{end_date or 'all'}.csv"
        df.to_csv(filename, index=False)
        return os.path.abspath(filename)


_stock_basic_df = None


def get_stock_basic():
    global _stock_basic_df
    if _stock_basic_df is None:
        pro = get_tushare_pro()
        _stock_basic_df = pro.stock_basic(
            exchange="", list_status="L", fields="ts_code,name,exchange,market,industry"
        )
    return _stock_basic_df


class TushareStockCodeLookup(Tool):
    name = "tushare_stock_code_lookup"
    description = """
    用于根据股票名称（支持模糊匹配，如'贵州茅台'、'茅台'）查询Tushare对应的股票代码"""
    inputs = {
        "stock_name": {
            "type": "string",
            "description": "股票名称或名称的一部分，如 '贵州茅台'、'平安银行'",
        }
    }
    output_type = "string"

    def forward(self, stock_name: str) -> str:
        df = get_stock_basic()
        if df is None or df.empty:
            return "获取股票列表失败，无法查询"
        exact = df[df["name"] == stock_name]
        if len(exact) == 1:
            row = exact.iloc[0]
            return f"{row['name']}的Tushare股票代码为：{row['ts_code']}（{row['exchange']}，{row['industry']}）"
        matched = df[df["name"].str.contains(stock_name, na=False)]
        if len(matched) == 1:
            row = matched.iloc[0]
            return f"{row['name']}的Tushare股票代码为：{row['ts_code']}（{row['exchange']}，{row['industry']}）"
        if len(matched) > 1:
            lines = [
                f"找到{len(matched)}个与'{stock_name}'相关的股票："
            ]
            for _, row in matched.head(20).iterrows():
                lines.append(
                    f"{row['name']}: {row['ts_code']}（{row['exchange']}，{row['industry']}）"
                )
            if len(matched) > 20:
                lines.append(f"... 等共{len(matched)}个，请使用更精确的名称重新查询")
            return "\n".join(lines)
        return f"未找到与'{stock_name}'相关的股票，请检查名称是否正确"


class BoChaWebSearch(Tool):
    name = "bocha_web_search"
    description = """
    用于根据搜索词调用Bocha Web Search API 进行联网搜索"""
    inputs = {
        "query": {
            "type": "string",
            "description": "要搜索的查询",
        }
    }
    output_type = "string"

    def forward(self, query: str) -> str:
        data = {"query": query, "summary": True, "count": 10, "page": 1}
        endpoint = "https://api.bochaai.com/v1/web-search"
        API_KEY = os.getenv("BOCHA_API_KEY")
        headers = {
            "Authorization": f"Bearer {API_KEY}",
            "Content-Type": "application/json",
        }

        response = requests.post(endpoint, headers=headers, data=json.dumps(data))
        response.raise_for_status()
        search_ret = response.json()
        return self.bocha_for_list(search_ret)

    def bocha_for_list(self, search_ret: dict):
        data = search_ret["data"]
        pages = data["webPages"]["value"]
        ret = []
        for page in pages:
            ret.append(
                {"title": page["name"], "summary": page["summary"], "url": page["url"]}
            )
        return ret