config.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
配置管理模块
加载环境变量和配置文件
"""

import os
from pathlib import Path
from dotenv import load_dotenv
from typing import Dict, Any

class ConfigManager:
    """配置管理类"""
    
    def __init__(self, env_file: str = ".env"):
        """初始化配置管理器"""
        # 加载环境变量
        load_dotenv(env_file)
        
        # 项目根目录
        self.root_dir = Path(__file__).parent.parent
        
        # 数据存储路径
        self.data_dir = self.root_dir / "data"
        self.data_dir.mkdir(exist_ok=True)
        
        # 结果存储路径
        self.results_dir = self.root_dir / "results"
        self.results_dir.mkdir(exist_ok=True)
        
        # 日志存储路径
        self.log_dir = self.root_dir / "logs"
        self.log_dir.mkdir(exist_ok=True)
        
        # 初始化配置
        self._init_config()
    
    def _init_config(self):
        """初始化配置"""
        # Tushare配置
        self.tushare_token = self._get_env("TUSHARE_TOKEN")
        
        # 火山引擎API配置
        self.volcengine_api_key = self._get_env("VOLCENGINE_API_KEY")
        self.volcengine_api_base = self._get_env("VOLCENGINE_API_BASE", "https://ark.cn-beijing.volces.com/api/v3/chat/completions")
        self.volcengine_model = self._get_env("VOLCENGINE_MODEL", "doubao-pro-32k")
        
        # 数据库配置
        self.db_path = self._get_env("DB_PATH", str(self.data_dir / "stock_data.db"))
        
        # 日志配置
        self.log_path = self._get_env("LOG_PATH", str(self.log_dir / "quant_trading.log"))
        
        # 选股参数
        self.selection_time = self._get_env("SELECTION_TIME", "14:30")
        self.stock_selection_count = int(self._get_env("STOCK_SELECTION_COUNT", 10))
        self.min_turnover_rate = float(self._get_env("MIN_TURNOVER_RATE", 1.0))
        self.max_stock_count = int(self._get_env("MAX_STOCK_COUNT", 3000))
        
        # 回测参数
        self.backtest_start_date = self._get_env("BACKTEST_START_DATE", "20200101")
        self.backtest_end_date = self._get_env("BACKTEST_END_DATE", "20241231")
        
        # 大模型参数
        self.temperature = float(self._get_env("TEMPERATURE", 0.1))
        self.top_p = float(self._get_env("TOP_P", 0.9))
        self.max_tokens = int(self._get_env("MAX_TOKENS", 8192))
        
        # 模型有效性阈值
        self.accuracy_threshold = float(self._get_env("ACCURACY_THRESHOLD", 0.5))
        self.profit_loss_threshold = float(self._get_env("PROFIT_LOSS_THRESHOLD", 2.0))
    
    def _get_env(self, key: str, default: str = None) -> str:
        """获取环境变量值"""
        value = os.getenv(key, default)
        if not value and not default:
            raise EnvironmentError(f"缺少环境变量：{key}")
        return value
    
    def get_config(self) -> Dict[str, Any]:
        """获取所有配置"""
        return {
            "tushare": {
                "token": self.tushare_token
            },
            "volcengine": {
                "api_key": self.volcengine_api_key,
                "api_base": self.volcengine_api_base,
                "model": self.volcengine_model
            },
            "database": {
                "path": self.db_path
            },
            "logging": {
                "path": self.log_path
            },
            "selection": {
                "time": self.selection_time,
                "count": self.stock_selection_count,
                "min_turnover_rate": self.min_turnover_rate,
                "max_stock_count": self.max_stock_count
            },
            "backtest": {
                "start_date": self.backtest_start_date,
                "end_date": self.backtest_end_date
            },
            "model": {
                "temperature": self.temperature,
                "top_p": self.top_p,
                "max_tokens": self.max_tokens
            },
            "validation": {
                "accuracy_threshold": self.accuracy_threshold,
                "profit_loss_threshold": self.profit_loss_threshold
            }
        }
    
    def print_config(self):
        """打印配置信息"""
        config = self.get_config()
        
        print("=" * 80)
        print("配置信息")
        print("=" * 80)
        
        for section, params in config.items():
            print(f"{section.upper()}")
            print("-" * len(section))
            
            for key, value in params.items():
                if key == "api_key" or key == "token":
                    print(f"{key}: {value[:5]}...")
                else:
                    print(f"{key}: {value}")
                    
            print()
        
        print("=" * 80)
    
    def save_config(self, config: Dict[str, Any], file_path: str = None):
        """保存配置到文件"""
        if not file_path:
            file_path = str(self.root_dir / "config" / "config.json")
        
        try:
            import json
            with open(file_path, 'w', encoding='utf-8') as f:
                json.dump(config, f, ensure_ascii=False, indent=2)
            
            return True
            
        except Exception as e:
            raise Exception(f"配置保存失败：{e}")
    
    def load_config(self, file_path: str = None) -> Dict[str, Any]:
        """从文件加载配置"""
        if not file_path:
            file_path = str(self.root_dir / "config" / "config.json")
        
        try:
            import json
            with open(file_path, 'r', encoding='utf-8') as f:
                return json.load(f)
                
        except FileNotFoundError:
            raise Exception("配置文件未找到")
        except Exception as e:
            raise Exception(f"配置加载失败：{e}")

# 全局配置实例
config_manager = ConfigManager()