开源自动化数据分析工具——data-copilot源码分析

泠临

已于 2024-08-08 14:39:23 修改

阅读量1.4k

点赞数 32

文章标签：自动化数据分析人工智能

于 2024-08-08 14:17:36 首次发布

本文链接：https://blog.csdn.net/qq_52159267/article/details/141021477

版权

本文对开源数据分析工具data-copilot源代码中的关键部分，create_tool/Atomic_api_json.py进行详细分析。这里是原论文地址与github项目地址。

创建工具部分create_tool/Atomic_api_json.py

这段代码的主要目的是使用 tushare 库获取各种金融数据，并将这些数据的描述、示例输入和输出保存到一个 JSON 文件中。以下是对代码的详细解释：

导入必要的库

import tushare as ts
import matplotlib.pyplot as plt
import pandas as pd
import os
import random
from matplotlib.ticker import MaxNLocator
import matplotlib.font_manager as fm
from prettytable import PrettyTable
from blessed import Terminal
import time
from datetime import datetime, timedelta
import numpy as np
import mplfinance as mpf
from prettytable import PrettyTable
from typing import Optional
import matplotlib.font_manager as fm
from matplotlib.lines import Line2D
from typing import Union, Any
import json

这些库提供了数据获取、数据处理、绘图、文件操作等功能。

获取 Tushare API 的 token 并初始化

token = os.getenv("TUSHARE_TOKEN")
pro = ts.pro_api(token)
all_atomic_api = {}

从环境变量中获取 Tushare 的 token，并初始化 pro 对象以便后续调用 API。

定义并获取各种金融数据

代码分多个部分，对于不同类型的金融数据，生成对应的获取数据的api，每个部分的结构大致相同：

获取数据
提取数据的第一行和最后一行
生成列名的中英文对照字典
创建描述数据的字典
将描述数据的字典添加到 all_atomic_api 中

以下是每个部分的详细解释：

财务指标数据

df = pro.fina_indicator(ts_code="600000.SH", start_date="20200104", end_date="20220104", fields=["ts_code", "end_date", "eps", "current_ratio", "quick_ratio", "inv_turn", "netprofit_margin", "grossprofit_margin", "roe", "roa", "roic", "debt_to_assets", "netprofit_yoy", "dt_netprofit_yoy"])
df_sample_str = df.iloc[[0, -1], :].to_string(header=False, index=False)

columns = df.columns.tolist()
columns_means = ['股票代码', '报告期', '每股收益', '流动比率', '速动比率', '存货周转率', '销售净利率', '销售毛利率', '净资产收益率', '总资产净利率', '投入资本回报率', '资产负债率', '净利润同比增长率', '扣非净利润同比增长率']
columns_dict = dict(zip(columns, columns_means))

atomic_api = {"func_desc": None, "Column_name": None, "example_input": None, "output_first_and_last_row": None}
atomic_api["func_desc"] = "财务指标"
atomic_api["Column_name"] = columns_dict
atomic_api["example_input"] = "pro.fina_indicator(ts_code = \"600000.SH\", start_date = \"20200104\", end_date = \"20220104\", fields=[\"ts_code\", \"end_date\", \"eps\", \"current_ratio\", \"quick_ratio\", \"inv_turn\", \"netprofit_margin\", \"grossprofit_margin\", \"roe\", \"roa\", \"roic\", \"debt_to_assets\", \"netprofit_yoy\", \"dt_netprofit_yoy\"])"
atomic_api["output_first_and_last_row"] = df_sample_str

all_atomic_api["pro.fina_indicator"] = atomic_api

这个部分获取了股票代码为 600000.SH 的财务指标数据，并将其描述信息保存到 all_atomic_api 字典中。让我们来看一下最终的结果，其实是生成了一个获取数据的api:

"pro.fina_indicator": {
        "func_desc": "财务指标",
        "Column_name": {
            "ts_code": "股票代码",
            "end_date": "报告期",
            "eps": "每股收益",
            "current_ratio": "流动比率",
            "quick_ratio": "速动比率",
            "inv_turn": "存货周转率",
            "netprofit_margin": "销售净利率",
            "grossprofit_margin": "销售毛利率",
            "roe": "净资产收益率",
            "roa": "总资产净利率",
            "roic": "投入资本回报率",
            "debt_to_assets": "资产负债率",
            "netprofit_yoy": "净利润同比增长率",
            "dt_netprofit_yoy": "扣非净利润同比增长率"
        },
        "example_input": "pro.fina_indicator(ts_code = \"600000.SH\",start_date = \"20200104\",end_date = \"20220104\",fields=[\"ts_code\",\"end_date\",\"eps\",\"current_ratio\",\"quick_ratio\",\"inv_turn\",\"netprofit_margin\",\"grossprofit_margin\",\"roe\",\"roa\",\"roic\",\"debt_to_assets\",\"netprofit_yoy\",\"dt_netprofit_yoy\"])",
        "output_first_and_last_row": "600000.SH 20211231 1.62 None None None 28.1524 None 8.1032 None None 91.6648 -9.1247 -9.26\n600000.SH 20200331 0.56 None None None 31.6289 None 3.0746 None None 92.0170  5.4803  5.23"
    }

获取上市公司基本信息

df = pro.stock_company(ts_code='600230.SH', fields=["ts_code", "exchange", "chairman", "manager", "secretary", "reg_capital", "setup_date", "province", "city", "introduction", "website", "email", "office", "employees", "main_business", "business_scope"])
df_sample_str = df.iloc[0, :].to_string(header=False, index=False)

columns = df.columns.tolist()
columns_means = ['股票代码', '交易所代码', '法人代表', '总经理', '董秘', '注册资本', '注册日期', '所在省份', '所在城市', '公司介绍', '公司主页', '电子邮件', '办公室地址', '员工人数', '主要业务及产品', '经营范围']
columns_dict = dict(zip(columns, columns_means))

atomic_api = {"func_desc": None, "Column_name": None, "example_input": None, "output_first_and_last_row": None}
atomic_api["func_desc"] = "获取上市公司基本信息如公司业务描述,员工人数等基本信息"
atomic_api["Column_name"] = columns_dict
atomic_api["example_input"] = "pro.stock_company(ts_code = '600230.SH', fields=[\"ts_code\", \"exchange\", \"chairman\", \"manager\", \"secretary\", \"reg_capital\", \"setup_date\", \"province\", \"city\", \"introduction\", \"website\", \"email\", \"office\", \"employees\", \"main_business\", \"business_scope\"])"
atomic_api["output_first_and_last_row"] = df_sample_str

all_atomic_api["pro.stock_company"] = atomic_api

这个部分获取了股票代码为 600230.SH 的上市公司基本信息，并将其描述信息保存到 all_atomic_api 字典中。

获取股票每日基本指标

df = pro.daily_basic(ts_code="600230.SH", start_date="20180726", end_date="20200726", fields=["ts_code", "trade_date", "turnover_rate", "turnover_rate_f", "volume_ratio", "pe_ttm", "pb", "ps_ttm", "dv_ttm", "total_share", "float_share", "free_share", "total_mv", "circ_mv"])
df_sample_str = df.iloc[[0, -1], :].to_string(header=False, index=False)

columns = df.columns.tolist()
columns_means = ['股票代码', '交易日期', '换手率(总股本)', '换手率(自由流通股本)', '量比', '市盈率(动态)', '市净率', '市销率(动态)', '股息率(动态)', '总股本', '流通股本', '自由流通股本', '总市值', '流通市值']
columns_dict = dict(zip(columns, columns_means))

atomic_api = {"func_desc": None, "Column_name": None, "example_input": None, "output_first_and_last_row": None}
atomic_api["func_desc"] = "获取股票每日基本指标例如换手率,市盈率市净率股息率等"
atomic_api["Column_name"] = columns_dict
atomic_api["example_input"] = "pro.daily_basic(ts_code = \"600230.SH\", start_date = \"20180726\", end_date = \"20200726\", fields=[\"ts_code\", \"trade_date\", \"turnover_rate\", \"turnover_rate_f\", \"volume_ratio\", \"pe_ttm\", \"pb\", \"ps_ttm\", \"dv_ttm\", \"total_share\", \"float_share\", \"free_share\", \"total_mv\", \"circ_mv\"])"
atomic_api["output_first_and_last_row"] = df_sample_str

all_atomic_api["pro.daily_basic"] = atomic_api

这个部分获取了股票代码为 600230.SH 的每日基本指标数据，并将其描述信息保存到 all_atomic_api 字典中。

获取股票的每日技术指标数据

df = pro.stk_factor(ts_code="600000.SH", start_date="20220520", end_date="20230520", fields=["ts_code", "trade_date", "close", "macd_dif", "macd_dea", "macd", "kdj_k", "kdj_d", "kdj_j", "rsi_6", "rsi_12", "rsi_24", "boll_upper", "boll_mid", "boll_lower", "cci"])
df_sample_str = df.iloc[[0, -1], :].to_string(header=False, index=False)

columns = df.columns.tolist()
columns_means = ['股票代码', '交易日期', '收盘价', 'MACD_DIF', 'MACD_DEA', 'MACD', 'KDJ_K', 'KDJ_D', 'KDJ_J', 'RSI_6', 'RSI_12', 'RSI_24', 'BOLL_UPPER', 'BOLL_MID', 'BOLL_LOWER', 'CCI']
columns_dict = dict(zip(columns, columns_means))

atomic_api = {"func_desc": None, "Column_name": None, "example_input": None, "output_first_and_last_row": None}
atomic_api["func_desc"] = "获取股票的每日的技术指标数据"
atomic_api["example_input"] = "pro.stk_factor(ts_code=stock_code, start_date=start_date, end_date=end_date, fields=[\"ts_code\", \"trade_date\", \"close\", \"macd_dif\", \"macd_dea\", \"macd\", \"kdj_k\", \"kdj_d\", \"kdj_j\", \"rsi_6\", \"rsi_12\", \"rsi_24\", \"boll_upper\", \"boll_mid\", \"boll_lower\", \"cci\"])"
atomic_api["output_first_and_last_row"] = df_sample_str
atomic_api["Column_name"] = str(columns_dict)

all_atomic_api["pro.stk_factor"] = atomic_api

这个部分获取了股票代码为 600000.SH 的每日技术指标数据，并将其描述信息保存到 all_atomic_api 字典中。

获取沪深港通资金每日的资金流向数据

df = pro.moneyflow_hsgt(start_date="20220101", end_date="20230101", fields=["trade_date","ggt_ss","ggt_sz","hgt","sgt","north_money","south_money"])
df_sample_str = df.iloc[[0, -1], :].to_string(header=False, index=False)
columns = df.columns.tolist()
columns_means = ['交易日期','港股通（上海）','港股通（深圳）','沪股通（百万元）','深股通（百万元）','北向资金（百万元）','南向资金（百万元）'  ]

columns_dict = dict(zip(columns, columns_means))

atomic_api = {"func_desc":None,"Column_name":None,"example_input":None,"output_first_and_last_row":None}

atomic_api["func_desc"] = "获取沪深港通资金每日的资金流向数据"
atomic_api["example_input"] = "pro.moneyflow_hsgt(start_date=\"20220101\", end_date=\"20230101\", fields=[\"trade_date\",\"ggt_ss\",\"ggt_sz\",\"hgt\",\"sgt\",\"north_money\",\"south_money\"])"
atomic_api["output_first_and_last_row"] = df_sample_str
atomic_api["Column_name"] = str(columns_dict)

all_atomic_api["pro.moneyflow_hsgt"] = atomic_api

这个部分获取了沪深港通资金每日的资金流向数据，并将其描述信息保存到 all_atomic_api 字典中。

获取申万行业指数的成分股信息

df = pro.index_member(index_code="850531.SI", fields=["index_code", "index_name", "con_code", "con_name", "in_date", "out_date", "is_new"])
df_sample_str = df.iloc[[0, -1], :].to_string(header=False, index=False)

columns = df.columns.tolist()
columns_means = ['指数代码', '指数名称', '成分股代码', '成分股名称', '纳入日期', '剔除日期', '是否最新']
columns_dict = dict(zip(columns, columns_means))

atomic_api = {"func_desc": None, "Column_name": None, "example_input": None, "output_first_and_last_row": None}
atomic_api["func_desc"] = "获取申万行业指数的成分股信息"
atomic_api["example_input"] = "pro.index_member(index_code= \"850531.SI \", fields=[\"index_code\", \"con_code\", \"in_date\", \"out_date\", \"is_new\", \"index_name\", \"con_name\"])"
atomic_api["output_first_and_last_row"] = df_sample_str
atomic_api["Column_name"] = str(columns_dict)

all_atomic_api["pro.index_member"] = atomic_api

这个部分获取了申万行业指数代码为 850531.SI 的成分股信息，并将其描述信息保存到 all_atomic_api 字典中。

获取申万一级(L1),二级(L2),三级(L3)的行业信息

df = pro.index_classify(level='L1', src='SW2021', fields=["index_code", "industry_name", "level"])
df_sample_str = df.iloc[[0, -1], :].to_string(header=False, index=False)

columns = df.columns.tolist()
columns_means = ['行业代码', '行业名称', '行业级别']
columns_dict = dict(zip(columns, columns_means))

atomic_api = {"func_desc": None, "Column_name": None, "example_input": None, "output_first_and_last_row": None}
atomic_api["func_desc"] = "获取申万一级(L1),二级(L2),三级(L3)的行业信息"
atomic_api["example_input"] = "pro.index_classify(level='L1', src='SW2021', filter=[\"index_code\", \"industry_name\", \"level\"])"
atomic_api["output_first_and_last_row"] = df_sample_str
atomic_api["Column_name"] = str(columns_dict)

all_atomic_api["pro.index_classify"] = atomic_api

这个部分获取了申万一级(L1),二级(L2),三级(L3)的行业信息，并将其描述信息保存到 all_atomic_api 字典中。

按照股票名称或者股票代码查询股票相关信息

df_stock_basic = pro.stock_basic(ts_code="", name="贵州茅台", fields=["ts_code", "name", "area", "industry", "market", "list_date"])
df_stock_basic_sample_str = df_stock_basic.iloc[0, :].to_string(header=False, index=False)

columns = df_stock_basic.columns.tolist()
columns_means = ['股票代码', '股票名称', '地域', '所属行业', '市场类型', '上市日期']
columns_dict = dict(zip(columns, columns_means))

atomic_api["func_desc"] = "按照股票名称或者股票代码查询股票相关信息"
atomic_api["example_input"] = "pro.stock_basic(ts_code=\"\", name=\"贵州茅台\", fields=[\"ts_code\", \"name\", \"area\", \"industry\", \"market\", \"list_date\"])"
atomic_api["output_first_and_last_row"] = df_stock_basic_sample_str
atomic_api["Column_name"] = str(columns_dict)

all_atomic_api["pro.stock_basic"] = atomic_api

这个部分获取了股票名称为“贵州茅台”的基本信息，并将其描述信息保存到 all_atomic_api 字典中。

获取股票每日的复权因子

df_adj_factor = pro.adj_factor(ts_code='000001.SZ', start_date='20180718', end_date='20190718', fields=["ts_code", "trade_date", "adj_factor"])
df_adj_factor_sample_str = df_adj_factor.iloc[[0, -1], :].to_string(header=False, index=False)

columns = df_adj_factor.columns.tolist()
columns_means = ['股票代码', '交易日期', '复权因子']
columns_dict = dict(zip(columns, columns_means))

atomic_api["func_desc"] = "获取股票每日的复权因子"
atomic_api["example_input"] = "pro.adj_factor(ts_code='000001.SZ', start_date='20180718', end_date='20190718', fields=[\"ts_code\", \"trade_date\", \"adj_factor\"])"
atomic_api["output_first_and_last_row"] = df_adj_factor_sample_str
atomic_api["Column_name"] = str(columns_dict)

all_atomic_api["pro.adj_factor"] = atomic_api

这个部分获取了股票代码为 000001.SZ 的每日复权因子，并将其描述信息保存到 all_atomic_api 字典中。

获取股票每日的行情数据

df_daily = pro.daily(ts_code='000001.SZ', start_date='20180701', end_date='20180718')
df_daily_sample_str = df_daily.iloc[[0, -1], :].to_string(header=False, index=False)

columns = df_daily.columns.tolist()
columns_means = ['股票代码', '交易日期', '开盘价', '最高价', '最低价', '收盘价', '昨收价', '涨跌额', '涨跌幅%', '成交量', '成交额']
columns_dict = dict(zip(columns, columns_means))

atomic_api["func_desc"] = "获取股票每日的行情数据"
atomic_api["example_input"] = "pro.daily(ts_code='000001.SZ', start_date='20180701', end_date='20180718')"
atomic_api["output_first_and_last_row"] = df_daily_sample_str
atomic_api["Column_name"] = str(columns_dict)

all_atomic_api["pro.daily"] = atomic_api

这个部分获取了股票代码为 000001.SZ 的每日行情数据，并将其描述信息保存到 all_atomic_api 字典中。

获取股票每周的行情数据

df_weekly = pro.weekly(ts_code='000001.SZ', start_date='20180101', end_date='20181101')
df_weekly_sample_str = df_weekly.iloc[[0, -1], :].to_string(header=False, index=False)

columns = df_weekly.columns.tolist()
columns_means = ['股票代码', '交易日期', '周收盘价', '周开盘价', '周最高价', '周最低价', '上一周收盘价', '周涨跌额', '周涨跌幅%', '周成交量', '周成交额']
columns_dict = dict(zip(columns, columns_means))

atomic_api["func_desc"] = "获取股票每周的行情数据"
atomic_api["example_input"] = "pro.weekly(ts_code='000001.SZ', start_date='20180101', end_date='20181101')"
atomic_api["output_first_and_last_row"] = df_weekly_sample_str
atomic_api["Column_name"] = str(columns_dict)

all_atomic_api["pro.weekly"] = atomic_api

这个部分获取了股票代码为 000001.SZ 的每周行情数据，并将其描述信息保存到 all_atomic_api 字典中。

获取股票每月的行情数据

df_monthly = pro.monthly(ts_code='000001.SZ', start_date='20180101', end_date='20181101')
df_monthly_sample_str = df_monthly.iloc[[0, -1], :].to_string(header=False, index=False)

columns = df_monthly.columns.tolist()
columns_means = ['股票代码', '交易日期', '月收盘价', '月开盘价', '月最高价', '月最低价', '上一月收盘价', '月涨跌额', '月涨跌幅%', '月成交量', '月成交额']
columns_dict = dict(zip(columns, columns_means))

atomic_api["func_desc"] = "获取股票每月的行情数据"
atomic_api["example_input"] = "pro.monthly(ts_code='000001.SZ', start_date='20180101', end_date='20181101')"
atomic_api["output_first_and_last_row"] = df_monthly_sample_str
atomic_api["Column_name"] = str(columns_dict)

all_atomic_api["pro.monthly"] = atomic_api

这个部分获取了股票代码为 000001.SZ 的每月行情数据，并将其描述信息保存到 all_atomic_api 字典中。

获取指数每日的行情数据

df = pro.index_daily(ts_code='399300.SZ', start_date='20180101', end_date='20191010')
df_sample_str = df.iloc[[0, -1], :].to_string(header=False, index=False)

columns = df.columns.tolist()
columns_means = ['指数代码', '交易日期', '收盘点位', '开盘点位', '最高点位', '最低点位', '昨收盘点位', '涨跌点', '涨跌幅%', '成交量', '成交额']
columns_dict = dict(zip(columns, columns_means))

atomic_api["func_desc"] = "获取指数每日的行情数据"
atomic_api["example_input"] = "pro.index_daily(ts_code='399300.SZ', start_date='20180101', end_date='20191010')"
atomic_api["output_first_and_last_row"] = df_sample_str
atomic_api["Column_name"] = str(columns_dict)

all_atomic_api["pro.index_daily"] = atomic_api

这个部分获取了指数代码为 399300.SZ 的每日行情数据，并将其描述信息保存到 all_atomic_api 字典中。

获取指数每周的行情数据

df = pro.index_weekly(ts_code='399300.SZ', start_date='20180101', end_date='20191010')
df_sample_str = df.iloc[[0, -1], :].to_string(header=False, index=False)

columns = df.columns.tolist()
columns_means = ['指数代码', '交易日期', '周收盘点位', '周开盘点位', '周最高点位', '周最低点位', '上周收盘点位', '周涨跌点', '周涨跌幅%', '周成交量', '周成交额']
columns_dict = dict(zip(columns, columns_means))

atomic_api["func_desc"] = "获取指数每周的行情数据"
atomic_api["example_input"] = "pro.index_weekly(ts_code='399300.SZ', start_date='20180101', end_date='20191010')"
atomic_api["output_first_and_last_row"] = df_sample_str
atomic_api["Column_name"] = str(columns_dict)

all_atomic_api["pro.index_weekly"] = atomic_api

这个部分获取了指数代码为 399300.SZ 的每周行情数据，并将其描述信息保存到 all_atomic_api 字典中。

获取指数每月的行情数据

df = pro.index_monthly(ts_code='399300.SZ', start_date='20180101', end_date='20191010')
df_sample_str = df.iloc[[0, -1], :].to_string(header=False, index=False)

columns = df.columns.tolist()
columns_means = ['指数代码', '交易日期', '月收盘点位', '月开盘点位', '月最高点位', '月最低点位', '上月收盘点位', '月涨跌点', '月涨跌幅%', '月成交量', '月成交额']
columns_dict = dict(zip(columns, columns_means))

atomic_api = {"func_desc": None, "Column_name": None, "example_input": None, "output_first_and_last_row": None}
atomic_api["func_desc"] = "获取指数每月的行情数据"
atomic_api["example_input"] = "pro.index_monthly(ts_code='399300.SZ', start_date='20180101', end_date='20191010')"
atomic_api["output_first_and_last_row"] = df_sample_str
atomic_api["Column_name"] = str(columns_dict)

all_atomic_api["pro.index_monthly"] = atomic_api

这个部分获取了指数代码为 399300.SZ 的每月行情数据，并将其描述信息保存到 all_atomic_api 字典中。

获取指数基本信息

df = pro.index_basic(name="沪深300", fields=["ts_code", "name"])
df_sample_str = df.iloc[[0, -1], :].to_string(header=False, index=False)

columns = df.columns.tolist()
columns_means = ['指数代码', '指数名称']
columns_dict = dict(zip(columns, columns_means))

atomic_api = {"func_desc": None, "Column_name": None, "example_input": None, "output_first_and_last_row": None}
atomic_api["func_desc"] = "获取指数基本信息"
atomic_api["example_input"] = "pro.index_basic(name='沪深300', fields=['ts_code', 'name'])"
atomic_api["output_first_and_last_row"] = df_sample_str
atomic_api["Column_name"] = str(columns_dict)

all_atomic_api["pro.index_basic"] = atomic_api

这个部分获取了指数名称为“沪深300”的基本信息，并将其描述信息保存到 all_atomic_api 字典中。

获取指数的成分股和成分股的权重

df = pro.index_weight(index_code="399300.SZ", start_date="20210101", end_date="20220101", fields=["index_code", "con_code", "trade_date", "weight"])
df_sample_str = df.iloc[[0, -1], :].to_string(header=False, index=False)

columns = df.columns.tolist()
columns_means = ['指数代码', '成分股代码', '交易日期', '成分股权重']
columns_dict = dict(zip(columns, columns_means))

atomic_api = {"func_desc": None, "Column_name": None, "example_input": None, "output_first_and_last_row": None}
atomic_api["func_desc"] = "获取指数的成分股和成分股的权重"
atomic_api["example_input"] = "pro.index_weight(index_code='399300.SZ', start_date='20210101', end_date='20220101', fields=['index_code', 'con_code', 'trade_date', 'weight'])"
atomic_api["output_first_and_last_row"] = df_sample_str
atomic_api["Column_name"] = str(columns_dict)

all_atomic_api["pro.index_weight"] = atomic_api

这个部分获取了指数代码为 399300.SZ 的成分股和成分股的权重，并将其描述信息保存到 all_atomic_api 字典中。

保存所有 API 描述信息到 JSON 文件

# 创建文件夹如果不存在
if not os.path.exists(""):
    os.mkdir("")

with open("all_atomic_api.json", "w") as f:
    json.dump(all_atomic_api, f, ensure_ascii=False, indent=4)

这段代码检查是否存在指定的文件夹，如果不存在则创建它。然后将 all_atomic_api 字典保存到一个名为 all_atomic_api.json 的文件中，确保 JSON 文件中的中文字符不会被转义，并且格式化输出（缩进为 4 个空格）。

总结

这段代码的主要功能是使用 Tushare 库获取各种金融数据，并将这些数据的描述、示例输入和输出保存到一个 JSON 文件中。通过这种方式，可以方便地查看和使用这些 API 的相关信息。

main.py中的run函数

这段代码分为多个步骤，下面将分步骤详细介绍这一段代码。

Step-1:Task select

第一步主要目的是通过调用语言模型来检测用户的意图，并根据意图生成任务计划

任务选择

获取当前时间

current_time = datetime.datetime.now()
formatted_time = current_time.strftime("%Y-%m-%d")
# 如果当前时间未超过下午3点，使用昨天的数据。
if current_time.hour < 15:
    formatted_time = (current_time - datetime.timedelta(days=1)).strftime("%Y-%m-%d")

这段代码获取当前时间，并将其格式化为 YYYY-MM-DD 的字符串。如果当前时间在下午3点之前，则使用前一天的日期。

打印意图检测阶段的提示

print('===============================Intent Detecting===========================================')

这行代码在控制台打印一条消息，表示意图检测阶段开始。

读取意图检测的提示词

with open('./prompt_lib/prompt_intent_detection.json', 'r') as f:
    prompt_task_dict = json.load(f)
prompt_intent_detection = ''
for key, value in prompt_task_dict.items():
    prompt_intent_detection = prompt_intent_detection + key + ": " + value + '\n\n'

这段代码从 prompt_lib/prompt_intent_detection.json 文件中读取意图检测的提示词，并将其转换为字符串格式。

构建意图检测的提示词

prompt_intent_detection = prompt_intent_detection + '\n\n' + 'Instruction:' + '今天的日期是' + formatted_time + ', ' + instruction + ' ###New Instruction: '

这段代码将当前日期和用户的指令添加到意图检测的提示词中。

发送意图检测请求

response = send_chat_request(model, prompt_intent_detection, openai_key=openai_key, api_base=api_base, engine=engine)

这段代码调用 send_chat_request 函数，向指定的语言模型发送意图检测请求，并获取响应。

处理意图检测的响应

new_instruction = response
print('new_instruction:', new_instruction)
output_text = output_text + '\n======Intent Detecting Stage=====\n\n'
output_text = output_text + new_instruction + '\n\n'

if add_to_queue is not None:
    add_to_queue(output_text)

这段代码处理意图检测的响应，将新的指令打印到控制台，并添加到输出文本中。如果存在 add_to_queue 回调函数，则将输出文本添加到队列中。

任务规划

打印任务规划阶段的提示

print('===============================Task Planing===========================================')
output_text = output_text + '=====Task Planing Stage=====\n\n'

这行代码在控制台打印一条消息，表示任务规划阶段开始，并将其添加到输出文本中。

读取任务规划的提示词

with open('./prompt_lib/prompt_task.json', 'r') as f:
    prompt_task_dict = json.load(f)
prompt_task = ''
for key, value in prompt_task_dict.items():
    prompt_task = prompt_task + key + ": " + value + '\n\n'

这段代码从 prompt_lib/prompt_task.json 文件中读取任务规划的提示词，并将其转换为字符串格式。

构建任务规划的提示词

prompt_task = prompt_task + '\n\n' + 'Instruction:' + new_instruction + ' ###Plan:'

这段代码将新的指令添加到任务规划的提示词中。

发送任务规划请求

response = send_chat_request(model, prompt_task, openai_key=openai_key, api_base=api_base, engine=engine)

这段代码调用 send_chat_request 函数，向指定的语言模型发送任务规划请求，并获取响应。

解析任务规划的响应

task_select = response
pattern = r"(task\d+=)(\{[^}]*\})"
matches = re.findall(pattern, task_select)
task_plan = {}
for task in matches:
    task_step, task_select = task
    task_select = task_select.replace("'", "\"")
    task_select = json.loads(task_select)
    task_name = list(task_select.keys())[0]
    task_instruction = list(task_select.values())[0]

    task_plan[task_name] = task_instruction

这段代码使用正则表达式解析任务规划的响应，将每个任务步骤和对应的任务指令存储在 task_plan 字典中。

打印和记录任务规划

for key, value in task_plan.items():
    print(key, ':', value)
    output_text = output_text + key + ': ' + str(value) + '\n'

output_text = output_text + '\n'
if add_to_queue is not None:
    add_to_queue(output_text)

这段代码将任务规划的每个步骤和指令打印到控制台，并添加到输出文本中。如果存在 add_to_queue 回调函数，则将输出文本添加到队列中。

总结

这段代码的主要功能是通过调用语言模型来检测用户的意图，并根据意图生成任务计划。它包括以下几个步骤：

获取当前时间并格式化。
读取意图检测的提示词，并构建完整的提示词。
发送意图检测请求，并处理响应。
读取任务规划的提示词，并构建完整的提示词。
发送任务规划请求，并解析响应。
打印和记录任务规划的结果。

通过这些步骤，代码能够自动检测用户的意图，并生成相应的任务计划，为后续的工具选择和使用阶段做准备。

Step-2:Tool select and use

这段代码的主要目的是通过调用语言模型来选择和使用工具，以完成特定任务。以下是代码的详细解释：

打印工具选择和使用阶段的提示

print('===============================Tool select and using Stage===========================================')
output_text = output_text + '======Tool select and using Stage======\n\n'

这行代码在控制台打印一条消息，表示工具选择和使用阶段开始，并将其添加到输出文本中。

读取任务选择的 JSON 文件名

task_name = list(task_plan.keys())[0].split('_task')[0]
task_instruction = list(task_plan.values())[0]

这段代码从 task_plan 字典中获取第一个任务的名称和指令。

构建工具库和提示文件路径

tool_lib = './tool_lib/' + 'tool_' + task_name + '.json'
tool_prompt = './prompt_lib/' + 'prompt_' + task_name + '.json'
prompt_flat = load_tool_and_prompt(tool_lib, tool_prompt)
prompt_flat = prompt_flat + '\n\n' + 'Instruction :' + task_instruction + ' ###Function Call'

这段代码构建工具库和提示文件的路径，并调用 load_tool_and_prompt 函数加载工具库和提示文件。然后将任务指令添加到提示词中。

发送工具选择和使用请求

response = send_chat_request(model, prompt_flat, openai_key=openai_key, api_base=api_base, engine=engine)

这段代码调用 send_chat_request 函数，向指定的语言模型发送工具选择和使用请求，并获取响应。

解析工具选择和使用的响应

if '###' in response:
    call_steps, _ = response.split('###')
else:
    call_steps = response
pattern = r"(step\d+=)(\{[^}]*\})"
matches = re.findall(pattern, call_steps)

这段代码检查响应中是否包含 ###，并根据情况分割响应。然后使用正则表达式解析响应中的步骤和内容。

初始化结果缓冲区

result_buffer = {}  # 存储格式如下：{'result1': (000001.SH, '中国平安的股票代码'), 'result2': (df2, '中国平安从1月到6月的股票数据')}.
output_buffer = []  # 存储变量名 [result5, result6]，将作为最终输出传递给下一个任务。

这段代码初始化结果缓冲区 result_buffer 和输出缓冲区 output_buffer。

处理每个匹配的步骤

for match in matches:
    step, content = match
    content = content.replace("'", "\"")  # 将单引号替换为双引号。
    print('==================')
    print("\n\nstep:", step)
    print('content:', content)
    call_dict = json.loads(content)
    print('It has parallel steps:', len(call_dict) / 4)
    output_text = output_text + step + ': ' + str(call_dict) + '\n\n'

    # 使用多线程并行执行以下代码。
    with concurrent.futures.ThreadPoolExecutor() as executor:
        # 提交任务到线程池
        futures = {executor.submit(parse_and_exe, call_dict, result_buffer, str(parallel_step))
                   for parallel_step in range(1, int(len(call_dict) / 4) + 1)}

        # 收集结果
        for idx, future in enumerate(concurrent.futures.as_completed(futures)):
            try:
                result = future.result()
                print('parallel step:', idx + 1)
            except Exception as exc:
                print(f'Generated an exception: {exc}')

    if step == matches[-1][0]:
        # 当前任务的最后一步。保存最终步骤的输出。
        for parallel_step in range(1, int(len(call_dict) / 4) + 1):
            output_buffer.append(call_dict['output' + str(parallel_step)])
output_text = output_text + '\n'
if add_to_queue is not None:
    add_to_queue(output_text)

这段代码处理每个匹配的步骤，将单引号替换为双引号，并将其转换为字典格式。然后使用多线程并行执行每个步骤，并收集结果。如果当前步骤是最后一步，则将输出添加到输出缓冲区 output_buffer 中。

总结

这段代码的主要功能是通过调用语言模型来选择和使用工具，以完成特定任务。它包括以下几个步骤：

打印工具选择和使用阶段的提示。
读取任务选择的 JSON 文件名。
构建工具库和提示文件路径，并加载工具库和提示文件。
发送工具选择和使用请求，并获取响应。
解析工具选择和使用的响应。
初始化结果缓冲区和输出缓冲区。
处理每个匹配的步骤，并使用多线程并行执行。

通过这些步骤，代码能够自动选择和使用合适的工具，以完成特定任务，并为后续的可视化阶段做准备。

Step-3:visualization

打印可视化阶段的提示

print('===============================Visualization Stage===========================================')
output_text = output_text + '======Visualization Stage====\n\n'

这行代码在控制台打印一条消息，表示可视化阶段开始，并将其添加到输出文本中。

获取任务名称和指令

task_name = list(task_plan.keys())[1].split('_task')[0]  # visualization_task
task_instruction = list(task_plan.values())[1]  # ''

这段代码从 task_plan 字典中获取第二个任务的名称和指令。

构建工具库和提示文件路径

tool_lib = './tool_lib/' + 'tool_' + task_name + '.json'
tool_prompt = './prompt_lib/' + 'prompt_' + task_name + '.json'

这段代码构建工具库和提示文件的路径。

初始化结果缓冲区

result_buffer_viz = {}
Previous_result = {}
for output_name in output_buffer:
    rename = 'input' + str(output_buffer.index(output_name) + 1)
    Previous_result[rename] = result_buffer[output_name][1]
    result_buffer_viz[rename] = result_buffer[output_name]

这段代码初始化结果缓冲区 result_buffer_viz 和前一个任务的结果 Previous_result。它将 output_buffer 中的每个输出重命名为 input1, input2 等，并将其添加到 result_buffer_viz 和 Previous_result 中。

加载工具和提示

prompt_flat = load_tool_and_prompt(tool_lib, tool_prompt)
prompt_flat = prompt_flat + '\n\n' + 'Instruction: ' + task_instruction + ', Previous_result: ' + str(Previous_result) + ' ###Function Call'

这段代码调用 load_tool_and_prompt 函数加载工具库和提示文件，并将任务指令和前一个任务的结果添加到提示词中。

发送可视化请求

response = send_chat_request(model, prompt_flat, openai_key=openai_key, api_base=api_base, engine=engine)

这段代码调用 send_chat_request 函数，向指定的语言模型发送可视化请求，并获取响应。

解析可视化的响应

if '###' in response:
    call_steps, _ = response.split('###')
else:
    call_steps = response
pattern = r"(step\d+=)(\{[^}]*\})"
matches = re.findall(pattern, call_steps)
for match in matches:
    step, content = match
    content = content.replace("'", "\"")  # 将单引号替换为双引号。
    print('==================')
    print("\n\nstep:", step)
    print('content:', content)
    call_dict = json.loads(content)
    print('It has parallel steps:', len(call_dict) / 4)
    result_buffer_viz = parse_and_exe(call_dict, result_buffer_viz, parallel_step='')
    output_text = output_text + step + ': ' + str(call_dict) + '\n\n'

这段代码检查响应中是否包含 ###，并根据情况分割响应。然后使用正则表达式解析响应中的步骤和内容，并调用 parse_and_exe 函数执行每个步骤，将结果存储在 result_buffer_viz 中。

添加输出文本到队列

if add_to_queue is not None:
    add_to_queue(output_text)

这段代码检查是否存在 add_to_queue 回调函数，如果存在，则将输出文本添加到队列中。

处理最终输出

finally_output = list(result_buffer_viz.values())  # plt.Axes

df = pd.DataFrame()
str_out = output_text + 'Finally result: '
for ax in finally_output:
    if isinstance(ax[0], plt.Axes):  # 如果输出是 plt.Axes，显示它。
        plt.grid()
        str_out = str_out + ax[1] + ':' + 'plt.Axes' + '\n\n'
    elif isinstance(ax[0], pd.DataFrame):
        df = ax[0]
        str_out = str_out + ax[1] + ':' + 'pd.DataFrame' + '\n\n'
    else:
        str_out = str_out + str(ax[1]) + ':' + str(ax[0]) + '\n\n'

这段代码处理最终的输出结果。如果输出是 plt.Axes 类型，则显示图表；如果是 pd.DataFrame 类型，则存储数据框。最后，将结果添加到 str_out 字符串中。

总结阶段

print('===============================Summary Stage===========================================')
output_prompt = "请用第一人称总结一下整个任务规划和解决过程,并且输出结果,用[Task]表示每个规划任务,用\{function\}表示每个任务里调用的函数." + \
                "示例1:###我用将您的问题拆分成两个任务,首先第一个任务[stock_task],我依次获取五粮液和贵州茅台从2013年5月20日到2023年5月20日的净资产回报率roe的时序数据. \n然后第二个任务[visualization_task],我用折线图绘制五粮液和贵州茅台从2013年5月20日到2023年5月20日的净资产回报率,并计算它们的平均值和中位数. \n\n在第一个任务中我分别使用了2个工具函数\{get_stock_code\},\{get_Financial_data_from_time_range\}获取到两只股票的roe数据,在第二个任务里我们使用折线图\{plot_stock_data\}工具函数来绘制他们的roe十年走势,最后并计算了两只股票十年ROE的中位数\{output_median_col\}和均值\{output_mean_col\}.\n\n最后贵州茅台的ROE的均值和中位数是\{\},{},五粮液的ROE的均值和中位数是\{\},\{\}###" + \
                "示例2:###我用将您的问题拆分成两个任务,首先第一个任务[stock_task],我依次获取20230101到20230520这段时间北向资金每日净流入和每日累计流入时序数据,第二个任务是[visualization_task],因此我在同一张图里同时绘制北向资金20230101到20230520的每日净流入柱状图和每日累计流入的折线图 \n\n为了完成第一个任务中我分别使用了2个工具函数\{get_north_south_money\},\{calculate_stock_index\}分别获取到北上资金的每日净流入量和每日的累计净流入量,第二个任务里我们使用折线图\{plot_stock_data\}绘制来两个指标的变化走势.\n\n最后我们给您提供了包含两个指标的折线图和数据表格." + \
                "示例3:###我用将您的问题拆分成两个任务,首先第一个任务[economic_task],我爬取了上市公司贵州茅台和其主营业务介绍信息. \n然后第二个任务[visualization_task],我用表格打印贵州茅台及其相关信息. \n\n在第一个任务中我分别使用了1个工具函数\{get_company_info\} 获取到贵州茅台的公司信息,在第二个任务里我们使用折线图\{print_save_table\}工具函数来输出表格.\n"
output_result = send_chat_request(model, output_prompt + str_out + '###', openai_key=openai_key, api_base=api_base, engine=engine)
print(output_result)

这段代码生成一个总结性的提示，并调用语言模型生成总结。然后将总结结果打印到控制台。

保存图表并返回结果

buf = BytesIO()
plt.savefig(buf, format='png')
buf.seek(0)

image = Image.open(buf)

return output_text, image, output_result, df

这段代码将图表保存为 PNG 格式的图片，并返回输出文本、图片、总结结果和数据框。

总结

这段代码的主要功能是通过调用语言模型来生成可视化图表，并总结整个任务的执行过程。它包括以下几个步骤：

打印可视化阶段的提示。
获取任务名称和指令。
构建工具库和提示文件路径，并加载工具库和提示文件。
初始化结果缓冲区和前一个任务的结果。
发送可视化请求，并解析响应。
处理最终输出结果。
生成总结性的提示，并调用语言模型生成总结。
保存图表并返回结果。

通过这些步骤，代码能够自动生成可视化图表，并总结整个任务的执行过程，为用户提供完整的解决方案。

tool.py

这个文件中就是各个功能api对应的实现函数，具体来说包括以下函数：

好的，以下是代码中各个函数的名字及其对应的功能：

1. **get_last_year_date**: 获取输入日期的前一年日期。
2. **get_adj_factor**: 获取股票的复权因子。
3. **get_stock_code**: 根据股票名称获取股票代码。
4. **get_stock_name_from_code**: 根据股票代码获取股票名称。
5. **get_stock_prices_data**: 获取指定股票在特定时间段内的每日、每周或每月价格数据。
6. **get_stock_technical_data**: 获取股票的每日技术指标数据。
7. **plot_stock_data**: 绘制股票数据的图表（折线图或柱状图）。
8. **query_fund_Manager**: 获取基金经理的信息。
9. **calculate_stock_index**: 计算股票的特定指标。
10. **rank_index_cross_section**: 根据给定指标对截面数据进行排序。
11. **get_company_info**: 获取公司的基本信息。
12. **get_Financial_data_from_time_range**: 获取指定股票在特定时间段内的财务数据。
13. **get_GDP_data**: 获取指定时间段内的 GDP 数据。
14. **get_cpi_ppi_currency_supply_data**: 获取指定时间段内的 CPI、PPI 或货币供应数据。
15. **predict_next_value**: 使用线性回归预测特定列的未来值。
16. **get_latest_new_from_web**: 获取最新的财经新闻。
17. **get_index_constituent**: 获取指定指数的成分股信息。
18. **is_fund**: 判断给定名称是股票还是基金。
19. **calculate_earning_between_two_time**: 计算指定股票或基金在两个日期之间的收益率。
20. **loop_rank**: 迭代应用给定函数并对结果进行排序。
21. **output_mean_median_col**: 计算指定列的均值和中位数。
22. **output_weighted_mean_col**: 计算指定列的加权平均值。
23. **get_index_data**: 获取指定指数的每日、每周或每月数据。
24. **get_north_south_money**: 获取指定时间段内的北向资金和南向资金流向数据。
25. **plot_k_line**: 绘制股票的 K 线图。
26. **cal_dt**: 计算两个时间点之间的百分比变化。
27. **query_fund_info**: 获取基金的基本信息。
28. **query_fund_data**: 获取基金的净值数据。
29. **query_fund_name_or_code**: 根据基金名称获取基金代码，或根据基金代码获取基金名称。
30. **print_save_table**: 打印并保存数据表格。
31. **merge_indicator_for_same_stock**: 合并同一股票的两个不同指标数据。
32. **select_value_by_column**: 选择数据框中特定列或特定值。

泠临

关注

32
点赞
踩
57

收藏

觉得还不错? 一键收藏
0
评论
复制链接

分享到 QQ

分享到新浪微博

扫一扫