stock_agent_2.py

import os
import sys
from datetime import datetime, timedelta
from typing import Optional
import dashscope
from qwen_agent.agents import Assistant
from qwen_agent.gui import WebUI
import pandas as pd
import numpy as np
from sqlalchemy import create_engine
from qwen_agent.tools.base import BaseTool, register_tool
import matplotlib.pyplot as plt
import time
from statsmodels.tsa.arima.model import ARIMA
import warnings

# 确保优先使用本地目录的 qwen_agent
current_dir = os.path.dirname(os.path.abspath(__file__))
if current_dir not in sys.path:
    sys.path.insert(0, current_dir)

# 解决中文显示问题
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'SimSun', 'Arial Unicode MS']
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 忽略警告
warnings.filterwarnings('ignore')

# 配置 DashScope
dashscope.api_key = os.getenv('DASHSCOPE_API_KEY', '')  # 从环境变量获取 API Key
dashscope.timeout = 120  # 设置超时时间为 120 秒（预测需要更长时间）

# ====== 股票助手 system prompt 和函数描述 ======
system_prompt = """我是股票助手，具备股票数据查询和价格预测功能。

## 数据表结构
CREATE TABLE stock_history (
    id INT AUTO_INCREMENT PRIMARY KEY,
    ts_code VARCHAR(20) NOT NULL COMMENT '股票代码',
    trade_date VARCHAR(10) NOT NULL COMMENT '交易日期 (格式：YYYY-MM-DD)',
    open DECIMAL(10, 2) COMMENT '开盘价',
    high DECIMAL(10, 2) COMMENT '最高价',
    low DECIMAL(10, 2) COMMENT '最低价',
    close DECIMAL(10, 2) COMMENT '收盘价',
    pre_close DECIMAL(10, 2) COMMENT '昨收价',
    change DECIMAL(10, 2) COMMENT '涨跌额',
    pct_chg DECIMAL(10, 4) COMMENT '涨跌幅 (%)',
    vol DECIMAL(15, 2) COMMENT '成交量 (手)',
    amount DECIMAL(15, 2) COMMENT '成交额 (千元)',
    stock_name VARCHAR(50) COMMENT '股票名称'
);

## 支持的股票
- 贵州茅台 (600519.SH)
- 五粮液 (000858.SZ)
- 国泰君安 (601211.SH)
- 中芯国际 (688981.SH)

## 可用工具
1. exc_sql: 执行 SQL 查询股票历史数据
2. arima_stock: 使用 ARIMA 模型预测股票未来 N 天的价格

## 预测功能说明
- 模型：ARIMA(5,1,5)
- 数据：获取截止到今天的前一年历史数据
- 输出：未来 N 天的价格预测（包含预测值和置信区间）

## 使用示例
1. 查询数据：贵州茅台 2024 年的收盘价
2. 预测价格：预测贵州茅台未来 7 天的价格
3. 预测价格：使用中芯国际股票代码，预测未来 10 天

每当工具返回结果时，你必须原样输出所有内容（包括表格和图片）。
"""

functions_desc = [
    {
        "name": "exc_sql",
        "description": "执行 SQL 查询股票历史数据",
        "parameters": {
            "type": "object",
            "properties": {
                "sql_input": {
                    "type": "string",
                    "description": "SQL 查询语句",
                }
            },
            "required": ["sql_input"],
        },
    },
    {
        "name": "arima_stock",
        "description": "使用 ARIMA 模型预测股票未来 N 天的价格",
        "parameters": {
            "type": "object",
            "properties": {
                "ts_code": {
                    "type": "string",
                    "description": "股票代码（必填），如：600519.SH",
                },
                "n_days": {
                    "type": "integer",
                    "description": "预测天数（可选，默认 7 天）",
                    "default": 7
                }
            },
            "required": ["ts_code"],
        },
    },
]

# ====== 数据库连接配置 ======
# 阿里云 MySQL 数据库连接
def get_db_engine():
    """获取数据库连接引擎"""
    # 请根据实际情况修改数据库连接信息
    return create_engine(
        'mysql+mysqlconnector://root:123456@localhost:3306/stock_db',
        connect_args={'connect_timeout': 10},
        pool_size=10,
        max_overflow=20
    )


# ====== 工具类实现 ======
@register_tool('exc_sql')
class ExcSQLTool(BaseTool):
    """SQL 查询工具"""
    description = '执行 SQL 查询股票历史数据，并自动可视化'
    parameters = [{
        'name': 'sql_input',
        'type': 'string',
        'description': 'SQL 查询语句',
        'required': True
    }]

    def call(self, params: str, **kwargs) -> str:
        import json
        args = json.loads(params)
        sql_input = args['sql_input']
        print("=============================")
        print('sql_input:' + sql_input)
        
        try:
            engine = get_db_engine()
            df = pd.read_sql(sql_input, engine)
            print("=============================")
            print('df:' + df)
            if df.empty:
                return "查询结果为空"
            
            md = df.head(10).to_markdown(index=False)
            
            # 生成图表
            save_dir = os.path.join(os.path.dirname(__file__), 'image_show')
            os.makedirs(save_dir, exist_ok=True)
            filename = f'bar_{int(time.time()*1000)}.png'
            save_path = os.path.join(save_dir, filename)
            
            generate_chart_png(df, save_path)
            img_path = os.path.join('image_show', filename)
            img_md = f'![图表]({img_path})'
            
            return f"{md}\n\n{img_md}"
        except Exception as e:
            return f"SQL 执行出错：{str(e)}"


@register_tool('arima_stock')
class ARIMAStockTool(BaseTool):
    """ARIMA 股票价格预测工具"""
    description = '使用 ARIMA(5,1,5) 模型预测股票未来 N 天的收盘价'
    parameters = [{
        'name': 'ts_code',
        'type': 'string',
        'description': '股票代码（必填），如：600519.SH',
        'required': True
    }, {
        'name': 'n_days',
        'type': 'integer',
        'description': '预测天数（可选，默认 7 天）',
        'required': False,
        'default': 7
    }]

    def call(self, params: str, **kwargs) -> str:
        import json
        args = json.loads(params)
        ts_code = args.get('ts_code')
        n_days = args.get('n_days', 7)
        
        if not ts_code:
            return "错误：股票代码 (ts_code) 是必填参数"
        
        try:
            # 1. 获取历史数据（截止到今天的前一年）
            end_date = datetime.now()
            start_date = end_date - timedelta(days=365)
            
            start_date_str = start_date.strftime('%Y-%m-%d')
            end_date_str = end_date.strftime('%Y-%m-%d')
            
            engine = get_db_engine()
            query = f"""
                SELECT trade_date, close 
                FROM stock_history 
                WHERE ts_code = '{ts_code}' 
                AND trade_date BETWEEN '{start_date_str}' AND '{end_date_str}'
                ORDER BY trade_date
            """
            
            df = pd.read_sql(query, engine)
            
            if df.empty or len(df) < 30:
                return f"错误：股票 {ts_code} 的历史数据不足，无法进行预测"
            
            # 2. 数据预处理
            df['trade_date'] = pd.to_datetime(df['trade_date'])
            df = df.set_index('trade_date')
            df['close'] = pd.to_numeric(df['close'], errors='coerce')
            df = df.dropna()
            
            if len(df) < 30:
                return f"错误：股票 {ts_code} 的有效数据不足，无法进行预测"
            
            # 3. 拟合 ARIMA(5,1,5) 模型
            print(f"正在拟合 ARIMA(5,1,5) 模型...")
            
            # 使用固定的 ARIMA(5,1,5) 参数
            model = ARIMA(
                df['close'],
                order=(5, 1, 5)  # ARIMA(5,1,5)
            )
            model_fit = model.fit()
            
            # 4. 预测未来 N 天
            print(f"正在预测未来 {n_days} 天的价格...")
            n_periods = n_days  # 预测 N 天
            forecast = model_fit.get_forecast(steps=n_periods)
            forecast_mean = forecast.predicted_mean
            conf_int = forecast.conf_int()
            
            # 5. 创建预测结果 DataFrame
            last_date = df.index[-1]
            future_dates = pd.date_range(
                start=last_date + pd.Timedelta(days=1),
                periods=n_periods,
                freq='B'  # 工作日频率
            )
            
            forecast_df = pd.DataFrame({
                '预测日期': future_dates.strftime('%Y-%m-%d'),
                '预测收盘价': forecast_mean.values,
                '置信区间下限': conf_int.iloc[:, 0].values,
                '置信区间上限': conf_int.iloc[:, 1].values
            })
            
            # 6. 生成表格
            md = forecast_df.to_markdown(index=False)
            
            # 7. 生成预测图表
            save_dir = os.path.join(os.path.dirname(__file__), 'image_show')
            os.makedirs(save_dir, exist_ok=True)
            filename = f'arima_forecast_{int(time.time()*1000)}.png'
            save_path = os.path.join(save_dir, filename)
            
            generate_forecast_chart(df, forecast_df, save_path)
            img_path = os.path.join('image_show', filename)
            img_md = f'![ARIMA 预测图]({img_path})'
            
            # 8. 返回结果
            result = f"## 股票 {ts_code} 未来 {n_days} 天价格预测（ARIMA 模型）\n\n"
            result += f"**历史数据范围**: {start_date_str} 至 {end_date_str}\n"
            result += f"**模型**: ARIMA(5,1,5)\n\n"
            result += f"### 预测结果\n{md}\n\n"
            result += f"### 预测趋势图\n{img_md}"
            
            return result
            
        except Exception as e:
            return f"预测出错：{str(e)}"


# ========== 可视化函数 ==========
def generate_chart_png(df_sql, save_path):
    """生成普通查询的图表"""
    columns = df_sql.columns
    if len(columns) < 2:
        return
    
    object_columns = df_sql.select_dtypes(include='O').columns.tolist()
    if len(columns) > 0 and columns[0] in object_columns:
        object_columns.remove(columns[0])
    
    num_columns = df_sql.select_dtypes(exclude='O').columns.tolist()
    
    use_line_chart = False
    if len(columns) > 0:
        first_col = columns[0]
        if 'date' in first_col.lower() or 'trade' in first_col.lower() or 'time' in first_col.lower():
            use_line_chart = True
        elif df_sql[first_col].dtype == 'object' and len(df_sql) > 1:
            try:
                pd.to_datetime(df_sql[first_col])
                use_line_chart = True
            except:
                pass
    
    x = np.arange(len(df_sql))
    
    if len(object_columns) > 0:
        pivot_df = df_sql.pivot_table(index=columns[0], columns=object_columns,
                                      values=num_columns, fill_value=0)
        fig, ax = plt.subplots(figsize=(10, 6))
        
        if use_line_chart:
            for col in pivot_df.columns:
                ax.plot(pivot_df.index, pivot_df[col], marker='o', linewidth=2, markersize=6, label=str(col))
            ax.set_title("股票数据走势")
        else:
            bottoms = None
            for col in pivot_df.columns:
                ax.bar(pivot_df.index, pivot_df[col], bottom=bottoms, label=str(col))
                if bottoms is None:
                    bottoms = pivot_df[col].copy()
                else:
                    bottoms += pivot_df[col]
            ax.set_title("股票数据统计")
    else:
        fig, ax = plt.subplots(figsize=(10, 6))
        
        if use_line_chart:
            for i, column in enumerate(columns[1:]):
                values = df_sql[column]
                ax.plot(x, values, marker='o', linewidth=2, markersize=6, label=column)
            ax.set_xticks(x)
            ax.set_xticklabels(df_sql[columns[0]], rotation=45, ha='right')
            ax.set_title("股票数据走势")
        else:
            bottom = np.zeros(len(df_sql))
            for column in columns[1:]:
                ax.bar(x, df_sql[column], bottom=bottom, label=column)
                bottom += df_sql[column]
            ax.set_xticks(x)
            ax.set_xticklabels(df_sql[columns[0]], rotation=45, ha='right')
            ax.set_title("股票数据统计")
    
    ax.legend()
    ax.set_xlabel(columns[0])
    ax.set_ylabel("数值")
    plt.tight_layout()
    plt.savefig(save_path)
    plt.close()


def generate_forecast_chart(historical_df, forecast_df, save_path):
    """生成 ARIMA 预测图表"""
    fig, ax = plt.subplots(figsize=(12, 6))
    
    # 绘制历史数据
    ax.plot(historical_df.index, historical_df['close'], 
            label='历史价格', linewidth=2, color='blue', alpha=0.7)
    
    # 绘制预测数据
    forecast_dates = pd.to_datetime(forecast_df['预测日期'])
    ax.plot(forecast_dates, forecast_df['预测收盘价'], 
            label='预测价格', linewidth=2, color='red', linestyle='--')
    
    # 绘制置信区间
    ax.fill_between(forecast_dates, 
                    forecast_df['置信区间下限'], 
                    forecast_df['置信区间上限'],
                    alpha=0.3, color='red', label='95% 置信区间')
    
    # 添加标题和标签
    ax.set_title('ARIMA 股票价格预测', fontsize=14, fontweight='bold')
    ax.set_xlabel('日期', fontsize=12)
    ax.set_ylabel('收盘价', fontsize=12)
    ax.legend(loc='best', fontsize=10)
    ax.grid(True, alpha=0.3)
    
    # 旋转 x 轴标签
    plt.xticks(rotation=45, ha='right')
    plt.tight_layout()
    
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    plt.close()


# ====== 初始化股票助手服务 ======
def init_agent_service():
    """初始化股票助手服务"""
    llm_cfg = {
        'model': 'qwen-max',
        'timeout': 120,
        'retry_count': 3,
        'api_type': 'qwen',
    }
    
    # MCP 工具配置
    tools = [{
        "mcpServers": {
            "tavily-mcp": {
                "command": "npx",
                "args": ["-y", "tavily-mcp@0.1.4"],
                "env": {
                    "TAVILY_API_KEY": "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"
                },
            }
        }
    }, 'exc_sql', 'arima_stock']
    
    try:
        bot = Assistant(
            llm=llm_cfg,
            name='股票助手',
            description='股票查询与价格预测',
            system_message=system_prompt,
            function_list=tools,
            files=['./faq.txt']
        )
        print("助手初始化成功！")
        return bot
    except Exception as e:
        print(f"助手初始化失败：{str(e)}")
        raise


def app_tui():
    """终端交互模式"""
    try:
        bot = init_agent_service()
        messages = []
        
        while True:
            try:
                query = input('user question: ')
                file = input('file url (press enter if no file): ').strip()
                
                if not query:
                    print('user question cannot be empty!')
                    continue
                
                if not file:
                    messages.append({'role': 'user', 'content': query})
                else:
                    messages.append({'role': 'user', 'content': [{'text': query}, {'file': file}]})
                
                print("正在处理您的请求...")
                response = []
                for response in bot.run(messages):
                    pass
                
                if response:
                    last_response = response[-1] if isinstance(response, list) else response
                    if isinstance(last_response, dict):
                        content = last_response.get('content', '')
                        if content:
                            print('\n' + '='*60)
                            print('AI 回复：')
                            print('='*60)
                            print(content)
                            print('='*60 + '\n')
                        else:
                            print('AI 回复：(无文本内容)')
                    else:
                        print('AI 回复:', response)
                else:
                    print('未收到 AI 回复')
                
                messages.extend(response)
            except Exception as e:
                print(f"处理请求时出错：{str(e)}")
                print("请重试或输入新的问题")
    except Exception as e:
        print(f"启动终端模式失败：{str(e)}")


def app_gui():
    """图形界面模式"""
    try:
        print("正在启动 Web 界面...")
        bot = init_agent_service()
        
        chatbot_config = {
            'prompt.suggestions': [
                '预测贵州茅台未来 7 天的价格',
                '使用中芯国际股票代码，预测未来 10 天的价格',
                '贵州茅台 2024 年的收盘价走势如何',
            ]
        }
        
        print("Web 界面准备就绪，正在启动服务...")
        app = WebUI(
            bot,
            chatbot_config=chatbot_config
        )
        app.run(
            server_name='127.0.0.1',
            server_port=7860,
            share=True
        )
    except Exception as e:
        print(f"启动 Web 界面失败：{str(e)}")
        print("请检查网络连接和 API Key 配置")


if __name__ == '__main__':
    # 运行模式选择
    app_tui()          # 终端交互模式
    # app_gui()          # 图形界面模式